import os
import sys
import numpy as np
import pandas as pd
from tqdm import tqdm
import torch
from torchvision import datasets, transforms

import utils
from model_class import MlpA
from constant import *


def main(n_runs = 10000):
    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
    bias_path_fp = os.path.join(SCRIPT_DIR, "model_para/fp_fc2_bias.pt")
    bias_path_tp = os.path.join(SCRIPT_DIR, "model_para/tp_fc2_bias.pt")
    bias_path_tp_moma = os.path.join(SCRIPT_DIR, "model_para/tp_fc2_moma_bias.pt")
    dtype = torch.float32

    input_size = 784
    hidden_size = 512
    output_size = 10
    batch_size = 2000
    mc_iter = int(2e4)
    dataset_name = "MNIST"

    seed = 42
    utils.set_seed(seed)

    cuda = torch.cuda.is_available()
    device = torch.device("cuda" if cuda else "cpu")
    num_workers = 0
    print(f"Cuda = {cuda} with num_workers = {num_workers}, system version = {sys.version}")

    METHODS = [
        ("EMP",    utils.propagate_moments,          True),
        ("MP",     utils.propagate_moments_wu2019,   False),
        ("PL-DNN", utils.propagate_moments_bibi2018, False),
    ]

    fp_bias_dict = torch.load(bias_path_fp, map_location=device)
    b_1_fp = fp_bias_dict["fc1.bias"]
    b_2_fp = fp_bias_dict["fc2.bias"]
    sigma_1_fp_acc = torch.full((hidden_size, 1), std_fp_1, device=device)
    sigma_2_fp_acc = torch.full((output_size, 1), std_fp_2, device=device)
    mean_1_fp_acc = torch.full((hidden_size, input_size), mean_fp_1, device=device)
    mean_2_fp_acc = torch.full((output_size, hidden_size), mean_fp_2, device=device)

    tp_bias_dict = torch.load(bias_path_tp, map_location=device)
    b_1_tp = tp_bias_dict["fc1.bias"]
    b_2_tp = tp_bias_dict["fc2.bias"]

    tp_moma_bias_dict = torch.load(bias_path_tp_moma, map_location=device)
    b_1_tp_moma = tp_moma_bias_dict["fc1.bias"]
    b_2_tp_moma = tp_moma_bias_dict["fc2.bias"]

    model_fp = MlpA(input_size, hidden_size, output_size)
    model_fp.load_state_dict(fp_bias_dict, strict=False)
    model_fp = utils.register_parameter_2layer_mlp(model_fp, mean_fp_1, var_fp_1, mean_fp_2, var_fp_2)
    model_fp = model_fp.to(device).eval()

    model_tp = MlpA(input_size, hidden_size, output_size)
    model_tp.load_state_dict(tp_bias_dict, strict=False)
    model_tp = utils.register_parameter_2layer_mlp(model_tp, mean_tp_1, var_tp_1, mean_tp_2, var_tp_2)
    model_tp = model_tp.to(device).eval()

    model_tp_moma = MlpA(input_size, hidden_size, output_size)
    model_tp_moma.load_state_dict(tp_moma_bias_dict, strict=False)
    model_tp_moma = utils.register_parameter_2layer_mlp(
        model_tp_moma, mean_tp_1_moma, var_tp_1_moma, mean_tp_2_moma, var_tp_2_moma
    )
    model_tp_moma = model_tp_moma.to(device).eval()

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
        transforms.Lambda(lambda x: x.view(-1)),
    ])
    test_set = datasets.MNIST(
        root=os.path.join(SCRIPT_DIR, "mnist_data"),
        train=False, download=True, transform=transform,
    )

    ratio_per_method = {name: {"fp": [], "tp": [], "tp_moma": []} for name, _, _ in METHODS}
    emp_uncertainty = {"fp": [], "tp": [], "tp_moma": []}

    for _ in tqdm(range(n_runs)):
        idx = np.random.randint(0, len(test_set))
        x_0, _ = test_set[idx]
        x_0 = x_0.to(device)

        mean_x = x_0.unsqueeze(0).unsqueeze(0)
        var_x = torch.zeros((1, 1, input_size), dtype=dtype, device=device)
        cov_x = torch.zeros((1, input_size, input_size), dtype=dtype, device=device)

        _, var_mc_fp_i, _ = utils.mc_stats_fp_batch(
            mc_iter, x_0, mean_1_fp_acc, sigma_1_fp_acc,
            mean_2_fp_acc, sigma_2_fp_acc, b_1_fp, b_2_fp,
            device=device, batch_size=batch_size,
        )
        _, var_mc_tp_i, _ = utils.mc_stats_tp_batch(
            mc_iter, x_0, p_1, q_1, p_2, q_2, b_1_tp, b_2_tp,
            device=device, batch_size=batch_size,
        )
        _, var_mc_tp_moma_i, _ = utils.mc_stats_tp_batch(
            mc_iter, x_0, p_1_moma, q_1_moma, p_2_moma, q_2_moma, b_1_tp_moma, b_2_tp_moma,
            device=device, batch_size=batch_size,
        )

        layers_fp = list(model_fp.modules())
        layers_tp = list(model_tp.modules())
        layers_tp_moma = list(model_tp_moma.modules())
        for name, mp_fn, supports_tp in METHODS:
            _, var_mp_fp_i, _ = mp_fn(mean_x, var_x, cov_x, layers_fp)
            ratio_per_method[name]["fp"].append(
                (var_mc_fp_i / var_mp_fp_i.squeeze()).cpu().numpy()
            )
            if name == "EMP":
                emp_uncertainty["fp"].append(float(var_mp_fp_i.sum().cpu()))

            if supports_tp:
                _, var_mp_tp_i, _ = mp_fn(mean_x, var_x, cov_x, layers_tp)
                ratio_per_method[name]["tp"].append(
                    (var_mc_tp_i / var_mp_tp_i.squeeze()).cpu().numpy()
                )
                _, var_mp_tp_moma_i, _ = mp_fn(mean_x, var_x, cov_x, layers_tp_moma)
                ratio_per_method[name]["tp_moma"].append(
                    (var_mc_tp_moma_i / var_mp_tp_moma_i.squeeze()).cpu().numpy()
                )
                if name == "EMP":
                    emp_uncertainty["tp"].append(float(var_mp_tp_i.sum().cpu()))
                    emp_uncertainty["tp_moma"].append(float(var_mp_tp_moma_i.sum().cpu()))

    csv_dir = os.path.join(SCRIPT_DIR, "output/")
    os.makedirs(csv_dir, exist_ok=True)

    def build_paper_table(precision):
        rows = []
        for o in range(output_size):
            row = {"Datasets": dataset_name if o == 0 else "", "o": o}
            for name, _, supports_tp in METHODS:
                if precision in ("tp", "tp_moma") and not supports_tp:
                    row[name] = "-"
                    continue
                arr = np.array(ratio_per_method[name][precision])
                row[name] = f"{arr[:, o].mean():.4f} ± {arr[:, o].std():.4f}"
            rows.append(row)
        return pd.DataFrame(rows)

    def print_table(title, df):
        print(f"\n=== {title} ===")
        print(df.to_string(index=False))

    table_fp = build_paper_table("fp")
    table_tp = build_paper_table("tp")
    table_tp_moma = build_paper_table("tp_moma")
    table_fp.to_csv(csv_dir + "table_fp.csv", index=False)
    table_tp.to_csv(csv_dir + "table_tp.csv", index=False)
    table_tp_moma.to_csv(csv_dir + "table_tp_moma.csv", index=False)

    uncertainty_table = pd.DataFrame([
        {"Model": "FP-MLP-2",      "Uncertainty": float(np.mean(emp_uncertainty["fp"]))},
        {"Model": "TP-MLP-2",      "Uncertainty": float(np.mean(emp_uncertainty["tp"]))},
        {"Model": "TP-MLP-2-MOMA", "Uncertainty": float(np.mean(emp_uncertainty["tp_moma"]))},
    ])
    uncertainty_table.to_csv(csv_dir + "uncertainty.csv", index=False)

    print_table("FP-MLP-2       (var_mc / var_mp,  precision=fp)", table_fp)
    print_table("TP-MLP-2       (var_mc / var_mp,  precision=tp)", table_tp)
    print_table("TP-MLP-2-MOMA  (var_mc / var_mp,  precision=tp_moma)", table_tp_moma)
    print_table("Uncertainty (EMP, sum of diag(cov), averaged over samples)", uncertainty_table)
    print(f"\nOutput written to {csv_dir}")


if __name__ == "__main__":
    # set n_runs = 10000 for final results, n_runs=5 for quick testing
    n_runs = 5
    main(n_runs)
