from typing import Dict
import pandas as pd
from investigation import storage_path, cwd
from investigation.cli.args import Analyze
from investigation.database.query import filter_entry, extrac_unique_config
from investigation.shell import case


def aggregate_analysis(args: Analyze, unknown_args: Dict, log_name: str = "history"):
    commit_dir = storage_path / args.commit
    storage_dir = commit_dir / "storage"
    cases_df, cases = extrac_unique_config(
        args.commit,
        batch_hash=args.batch,
        **unknown_args,
    )
    print(cases_df)
    # print(cases)
    target_cols = args.targets

    agg = args.aggregate

    print(f"Analyzing {len(cases)} cases on targets {target_cols} with aggregate method {agg}.")
    csv_files = [x for x in storage_dir.glob(f"**/{log_name}.csv") if x.exists() and x.stat().st_size > 0]
    # print(csv_files)
    df_dict = {
        csv_file.relative_to(commit_dir / "storage").parts[0]: pd.read_csv(csv_file)
        for csv_file in csv_files if csv_file.exists() and csv_file.stat().st_size > 0
    }
    print({k: v.shape for k, v in df_dict.items()})

    col_inst_dict: Dict[str, pd.DataFrame] = {}
    for target in target_cols:
        parts = [df_dict[part][target].rename(part) for part in df_dict.keys() if target in df_dict[part].columns]
        if len(parts) == 0:
            print(f"Warning: No data found for target {target}.")
            continue
        df = pd.concat(
            parts,
            # ignore_index=True,
            join="outer",
            axis=1,
        )
        # Remove rows with all NaN values except for the global_step column
        df = df.dropna(how="all", subset=df.columns.difference(['global_step']), axis=0)
        col_inst_dict[target] = df
    print({k: v.shape for k, v in col_inst_dict.items()})
    # Process aggregation
    agg_dict = {}
    match agg:
        case "none":
            for target in target_cols:
                agg_dict[target] = col_inst_dict[target]

        case "mean" | "sum" | "max" | "min":
            for target in target_cols:
                if (target not in col_inst_dict):
                    continue
                agg_dict[target] = col_inst_dict[target].agg(axis=0, func=agg)

        case _:
            print(f"Unknown aggregate method {agg}, please use 'none', 'mean', 'sum', 'max' or 'min'.")
            exit(1)
    # print({k: v.shape for k, v in agg_dict.items()})

    concatenated_df = pd.DataFrame(agg_dict)

    if args.seed_merge:

        # Merge rows with the same configuration but different seeds by averaging their values
        config_cols = [col for col in cases_df.columns if col != 'seed']
        cases_df_groups = cases_df.groupby(config_cols, as_index=False, dropna=False)
        index_group_map = {}
        for idx, (name, group) in enumerate(cases_df_groups):
            for case_idx in group.index:
                index_group_map[case_idx] = idx
        assert len(index_group_map) == len(
            cases_df), f"Index group map length {len(index_group_map)} does not match cases_df length {len(cases_df)}"
        cases_df['merged_index'] = cases_df.index.map(index_group_map)
        # config_df = cases_df.drop(columns=['seed']).drop_duplicates().set_index('merged_index')
        config_df = cases_df.groupby('merged_index', as_index=True, dropna=False).first().drop(columns=['seed'])
        counts = cases_df.groupby('merged_index', as_index=True, dropna=False).size().rename('seed_count')
        config_df = pd.concat([config_df, counts], axis=1)

        concatenated_df['merged_index'] = concatenated_df.index.map(index_group_map)
        concatenated_df_mean = concatenated_df.groupby('merged_index').agg('mean')
        # print(f"Merged aggregated DataFrame by averaging over seeds, new shape: {concatenated_df_mean.shape}")
        concatenated_df_std = concatenated_df.groupby('merged_index').agg('std')
        # print(f"Merged aggregated DataFrame by std over seeds, new shape: {concatenated_df_std.shape}")
        # Combine mean and std into multi-level columns
        concatenated_df_mean.rename(columns=lambda x: f"{x}_seedmean", inplace=True)
        concatenated_df_std.rename(columns=lambda x: f"{x}_seedstd", inplace=True)
        concatenated_df = pd.concat([concatenated_df_mean, concatenated_df_std], axis=1, join="inner")
        concatenated_df.sort_index(axis=1, inplace=True)
        cases_df = config_df
        # print(cases_df)
        # print(concatenated_df)
        # assert False

    print(f"Concatenated aggregated DataFrame:\n{concatenated_df}")
    full_df = pd.concat(
        [cases_df, concatenated_df],
        axis=1,
        join="inner",
    )
    full_df.index.name = cases_df.index.name or "job_name"
    print(f"Full DataFrame with case configurations:\n{full_df.to_string(index=True)}")

    match args.output_format:
        case "csv":
            if args.output_group != "":
                output_path = cwd / "analysis_outputs" / args.output_group
                output_path.parent.mkdir(parents=True, exist_ok=True)
                full_df.to_csv(output_path.with_suffix(".csv"), index=True)
            else:
                print(full_df.to_csv(index=True))
        case "latex":
            latex_table, content = export_latex_table(full_df)
            if args.output_group != "":
                output_path = cwd / "analysis_outputs" / args.output_group
                output_path.parent.mkdir(parents=True, exist_ok=True)
                with open(output_path.with_suffix(".tex"), "w") as f:
                    f.write(latex_table)
            else:
                print(latex_table)
        case _:
            print(full_df)


def export_latex_table(df: pd.DataFrame):
    content = [
        "\\begin{tabular}{%s}" % ("l" * len(df.columns)),
        " & ".join(df.columns) + " \\\\ \\hline",
        *[f"{' & '.join(map(str, row))} \\\\ \\hline" for row in df.values],
        "\\end{tabular}",
    ]
    return "\n".join(content), content
