main.py

# Copyright (c) 2026 CyberAgent AI Lab
# Author: Ayuto Tsutsumi (@Atotti) <aya172957@ayutaso.com>

import argparse
from pathlib import Path

import random
import numpy as np

from src.character import Character
from src.matrix import Matrix
from src.situation import Situation
from src.utils import update_id_map
from src.visualization import visualize_embeddings
from src.eval import eval_matrix

parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(dest="task", required=True)

generate = subparsers.add_parser("generate")
generate.add_argument("--num", type=int, default=1)
generate.add_argument("--output_dir", type=str, default="data/processed/matrix")
generate.add_argument("--model_name", type=str, default="Qwen/Qwen3-1.7B")
generate.add_argument("--temperature", type=float, default=0.7)
generate.add_argument("--style", type=str, default="vanilla")
generate.add_argument("--master", type=str, default="data/external/character-situation-master.json")

master_embedding = subparsers.add_parser("master_embedding")
master_embedding.add_argument("--input_path", type=str, default="data/processed/matrix/blocks.json")
master_embedding.add_argument("--output_dir", type=str, default="data/processed/matrix")
master_embedding.add_argument("--master", type=str, default="data/external/character-situation-master.json")

style_embedding = subparsers.add_parser("style_embedding")
style_embedding.add_argument("--input_path", type=str, default="data/processed/matrix/blocks.json")
style_embedding.add_argument("--output_dir", type=str, default="data/processed/matrix")
style_embedding.add_argument("--master", type=str, default="data/external/character-situation-master.json")

text_embedding = subparsers.add_parser("text_embedding")
text_embedding.add_argument("--input_path", type=str, default="data/processed/matrix/blocks.json")
text_embedding.add_argument("--output_dir", type=str, default="data/processed/matrix")
text_embedding.add_argument("--master", type=str, default="data/external/character-situation-master.json")
text_embedding.add_argument(
    "--id_map_path",
    type=str,
    default=None,
    help="Path to ID map JSON file. If not provided, it will be derived from the master data path.",
)

optimize_parser = subparsers.add_parser("optimize")
optimize_parser.add_argument("--input_path", type=str, default="data/processed/matrix/blocks.json")
optimize_parser.add_argument("--output_path", type=str, default="data/processed/matrix/optimized_matrix.csv")
optimize_parser.add_argument(
    "--mode",
    type=str,
    choices=[
        "powell",
        "random",
        "optuna",
        "optuna_sum",
        "optuna_local_global",
        "optuna_local_consistency_diversity",
        "optuna_local_style_content",
    ],
    default="optuna",
    help="Optimization method to use",
)
optimize_parser.add_argument(
    "--weights",
    type=float,
    nargs=4,
    default=[1.0, 1.0, 1.0, 1.0],
    help="Weights for Situation Diversity (SD), Situation Consistency (SC),\
Content Diversity (CD), and Content Consistency (CC)",
    metavar=("SD", "SC", "CD", "CC"),
)
optimize_parser.add_argument("--n_trials", type=int, default=50, help="Number of trials for optuna optimization")
optimize_parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
optimize_parser.add_argument("--master", type=str, default="data/external/character-situation-master.json")


visualize = subparsers.add_parser("visualize")
visualize.add_argument("--input_dir", type=str, default="data/processed/matrix")
visualize.add_argument("--output_dir", type=str, default="data/processed/matrix/visual")
visualize.add_argument("--master", type=str, default="data/external/character-situation-master.json")

evaluate = subparsers.add_parser("evaluate")
evaluate.add_argument(
    "--judge_model",
    type=str,
    default=None,
    help="Model name for the local judge. Defaults to 'Qwen/Qwen3-1.7B'.",
)

args = parser.parse_args()
print(f"=== Starting task: {args.task} ===")

if args.task == "generate":
    num = args.num
    output_dir = args.output_dir
    model_name = args.model_name
    temperature = args.temperature
    style = args.style
    master_path = Path(args.master)
    print(f"{num=}\n{model_name=}\n{temperature=}\n{style=}")
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # Derive id_map path from master data path
    id_map_path = Path("data/interim") / f"{master_path.stem}_id_map.json"

    # Ensure the ID map is up to date before loading
    print(f"Updating ID map at: {id_map_path}")
    try:
        update_id_map.update_id_mappings(master_path, id_map_path)
    except Exception as e:
        print(f"Error updating ID map: {e}")
        raise

    character_list = Character.load_from_master(str(master_path), str(id_map_path))
    situation_list = Situation.load_from_master(str(master_path), str(id_map_path))
    matrix = Matrix(character_list, situation_list)
    matrix.generate_matrix(num, model_name, temperature, style=style)
    matrix.save(output_dir)
    generation_config = Path(output_dir) / "generation.config"
    with generation_config.open("w", encoding="utf-8") as f:
        f.write(f"{num=}\n{model_name=}\n{temperature=}\n{style=}")

if args.task == "master_embedding":
    input_path = args.input_path
    output_dir = args.output_dir
    master_path = Path(args.master)
    id_map_path = Path("data/interim") / f"{master_path.stem}_id_map.json"

    blocks_json_path = Path(input_path)
    matrix = Matrix.load(
        blocks_json=str(blocks_json_path), master_data_path=str(master_path), id_map_path=str(id_map_path)
    )

    # Calculate embeddings for characters and situations
    for character in matrix.character_list:
        character.calc_style_ref_embedding()
        character.calc_background_embedding()
        print(character)
        character.save()

    for situation in matrix.situation_list:
        situation.calc_content_embedding()
        print(situation)
        situation.save()

if args.task == "style_embedding":
    input_path = args.input_path
    output_dir = args.output_dir
    master_path = Path(args.master)
    id_map_path = Path("data/interim") / f"{master_path.stem}_id_map.json"

    # Load matrix data
    blocks_json = Path(input_path)
    matrix = Matrix.load(blocks_json=str(blocks_json), master_data_path=str(master_path), id_map_path=str(id_map_path))

    matrix.calc_embeddings(mode="style_embedding")
    matrix.save(output_dir)

if args.task == "text_embedding":
    input_path = args.input_path
    output_dir = args.output_dir
    master_path = Path(args.master)
    id_map_path = args.id_map_path
    if id_map_path is None:
        id_map_path = Path("data/interim") / f"{master_path.stem}_id_map.json"

    # Load matrix data
    blocks_json = Path(input_path)
    matrix = Matrix.load(blocks_json=str(blocks_json), master_data_path=str(master_path), id_map_path=str(id_map_path))

    matrix.calc_embeddings(mode="text_embedding")
    matrix.save(output_dir)

if args.task == "compute_tensor":
    input_path = args.input_path
    output_path = args.output_path
    mode = args.mode
    weights = tuple(args.weights)
    n_trials = args.n_trials
    seed = args.seed

    # Note that fixing random seeds does not guarantee complete reproducibility when running in parallel.
    # For debugging or exact reproducibility, run in a single process.
    random.seed(seed)
    np.random.seed(seed)  # Legacy code might use numpy's global random seed

    # Load matrix data
    blocks_json = Path(input_path)
    master_path = Path(args.master)
    id_map_path = Path("data/interim") / f"{master_path.stem}_id_map.json"

    # TODO: Matrix.load is taking around 20 minutes as it requires loading all texts and block embeddings one by one.
    # Make it load the whole tensor and texts from a single file.
    # This will reduce the time for loading the whole matrix again, making the optimization much faster and convenient.
    matrix = Matrix.load(blocks_json=str(blocks_json), master_data_path=str(master_path), id_map_path=str(id_map_path))

    # Generate a tensor of shape (num_characters, num_situations, num_samples, embedding_dim) in numpy format and cache it

    tensor, situation_embeddings, character_embeddings = matrix.to_numpy()
    tensor_cache_path = Path("{input_path}_dialogue_embeddings.npy".format(input_path=input_path.replace(".json", "")))
    np.save(tensor_cache_path, tensor)
    print(f"Cached tensor to {tensor_cache_path} with shape {tensor.shape}")

    situation_cache_path = Path(
        "{input_path}_situation_embeddings.npy".format(input_path=input_path.replace(".json", ""))
    )
    np.save(situation_cache_path, situation_embeddings)
    print(f"Cached situation embeddings to {situation_cache_path} with shape {situation_embeddings.shape}")

    character_cache_path = Path(
        "{input_path}_character_embeddings.npy".format(input_path=input_path.replace(".json", ""))
    )
    np.save(character_cache_path, character_embeddings)
    print(f"Cached character embeddings to {character_cache_path} with shape {character_embeddings.shape}")

    exit(0)

    # Export optimized matrix to csv
    # mx, index_matrix, component_values = matrix.optimize(mode=mode, output=output_path, weights=weights, n_trials=n_trials)
    list_of_optimized_matrices = matrix.optimize(mode=mode, output=output_path, weights=weights, n_trials=n_trials)

    for i in range(len(list_of_optimized_matrices)):
        mx, index_matrix, component_values = list_of_optimized_matrices[i]

        optimization_config = Path(output_path.replace(".csv", f"_{i}.csv") + ".config")
        index_matrix_str = "\n".join([str(index) for index in index_matrix])

        # Format component values for config file
        component_values_str = "\n".join([f"{k}={v}" for k, v in component_values.items()])

        with optimization_config.open("w", encoding="utf-8") as f:
            f.write(
                f"{mode=}\n{weights=}\n\n{component_values_str}\n\
    \nindex_matrix=\n{index_matrix_str}"
            )

if args.task == "optimize":
    input_path = args.input_path
    output_path = args.output_path
    mode = args.mode
    weights = tuple(args.weights)
    n_trials = args.n_trials
    seed = args.seed

    # Note that fixing random seeds does not guarantee complete reproducibility when running in parallel.
    # For debugging or exact reproducibility, run in a single process.
    random.seed(seed)
    np.random.seed(seed)  # Legacy code might use numpy's global random seed

    # load dialogue embeddings tensor, situation embeddings, and character embeddings
    tensor_cache_path = Path("{input_path}_dialogue_embeddings.npy".format(input_path=input_path.replace(".json", "")))
    tensor = np.load(tensor_cache_path)
    print(f"Loaded tensor from {tensor_cache_path} with shape {tensor.shape}")

    situation_cache_path = Path(
        "{input_path}_situation_embeddings.npy".format(input_path=input_path.replace(".json", ""))
    )
    situation_embeddings = np.load(situation_cache_path)
    print(f"Loaded situation embeddings from {situation_cache_path} with shape {situation_embeddings.shape}")

    character_cache_path = Path(
        "{input_path}_character_embeddings.npy".format(input_path=input_path.replace(".json", ""))
    )
    character_embeddings = np.load(character_cache_path)
    print(f"Loaded character embeddings from {character_cache_path} with shape {character_embeddings.shape}")

    # TODO: Implement optimization logic here
    from src.optimization.jax import optimizer

    obj_name = mode.replace("optuna_", "")

    index_matrices, component_values = optimizer(
        dialogue_data=tensor,
        situation_data=situation_embeddings,
        character_data=character_embeddings,
        weights=weights,
        n_trials=n_trials,
        obj_name=obj_name,
    )
    for i in range(len(index_matrices)):
        index_matrix = index_matrices[i]
        component_value = component_values[i]

        optimization_config = Path(output_path.replace(".csv", f"_{obj_name}_optimized_{i}.csv") + ".config")
        index_matrix_str = "\n".join([str(index) for index in index_matrix])

        # Format component values for config file
        component_value_str = "\n".join([f"{k}={v}" for k, v in component_value.items()])

        with optimization_config.open("w", encoding="utf-8") as f:
            f.write(f"{mode=}\n{weights=}\n\n{component_value_str}\n\nindex_matrix=\n{index_matrix_str}")

        index_matrix_np = np.array(index_matrix)
        np.save(output_path.replace(".csv", f"_{obj_name}_optimized_{i}.npy"), index_matrix_np)


if args.task == "visualize":
    input_path = args.input_dir
    output_dir = args.output_dir
    master_path = Path(args.master)
    id_map_path = Path("data/interim") / f"{master_path.stem}_id_map.json"

    # Load matrix data
    blocks_json = Path(input_path) / "blocks.json"
    matrix = Matrix.load(blocks_json=str(blocks_json), master_data_path=str(master_path), id_map_path=str(id_map_path))

    visualize_embeddings(matrix, output_dir)


if args.task == "evaluate":
    import csv

    judge_model = args.judge_model

    # TODO: hard coded for the two methods (optimized and random). Make it more flexible.
    results = eval_matrix.evaluate_absolute(judge_model=judge_model)
    summary = results["summary"]
    points = results["points"]

    print("Evaluation Results:")
    print(f"Optimized matrices: {summary['optimized']}")
    print(f"Random matrices: {summary['random']}")

    # Save points to CSV
    output_path = Path("data/processed/evaluation_points.csv")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with output_path.open("w", newline="", encoding="utf-8") as csvfile:
        fieldnames = ["filename", "point"]
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(points)

    print(f"\nPoints saved to: {output_path}")


print(f"=== Task completed: {args.task} ===")