import sys
sys.path.append('./src')
import matplotlib
matplotlib.use('Agg')

from omegaconf import OmegaConf ### for yaml config parsing
import torch
import numpy as np 
import torch.optim as optim
from tqdm import tqdm
from torchinfo import summary
from pathlib import Path

from models import get_triple_photon_model_class 
from datasets import triplePhotonDataset 

### random seed for reproducibility
torch.manual_seed(0)
torch.cuda.manual_seed(0)
np.random.seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

conf = OmegaConf.load("Configs/train_3photon.yaml")

def prepare_output_folder(conf):
    from datetime import datetime
    date = datetime.now().strftime("%y%m%d") ## YYMMDD format
    # find the next index for experiment name
    exp_index = 0
    while True:
        exp_name = f'{date}_{conf.data.energy}keV_v{conf.model.version}_{exp_index:02d}'
        if not Path(f'Results/3ph/{exp_name}').exists():
            break
        exp_index += 1
    Path(f'Results/3ph/{exp_name}').mkdir(parents=True, exist_ok=True)
    Path(f'Results/3ph/{exp_name}/Models').mkdir(parents=True, exist_ok=True)
    Path(f'Results/3ph/{exp_name}/Plots').mkdir(parents=True, exist_ok=True)
    OmegaConf.save(conf, f'Results/3ph/{exp_name}/config.yaml')
    return exp_name

def get_loss_function(conf):
    if conf.loss.type == "three_point_set_loss_smooth_l1":
        def three_point_set_loss_smooth_l1(pred_xy, gt_xy):
            loss_fn = torch.nn.SmoothL1Loss(reduction='none', beta=conf.loss.smooth_l1_beta)
            
            p1, p2, p3 = pred_xy[:,0], pred_xy[:,1], pred_xy[:,2]
            g1, g2, g3 = gt_xy[:,0], gt_xy[:,1], gt_xy[:,2]
            
            # 1. 匹配阶段：计算全部 6 种排列的 L2 物理距离 (不参与梯度计算)
            with torch.no_grad():
                def calc_cost(ga, gb, gc):
                    return ((p1 - ga)**2).sum(dim=-1) + \
                           ((p2 - gb)**2).sum(dim=-1) + \
                           ((p3 - gc)**2).sum(dim=-1)
                
                # 枚举 3个光子的 6 种全排列 (Permutations)
                c0 = calc_cost(g1, g2, g3)
                c1 = calc_cost(g1, g3, g2)
                c2 = calc_cost(g2, g1, g3)
                c3 = calc_cost(g2, g3, g1)
                c4 = calc_cost(g3, g1, g2)
                c5 = calc_cost(g3, g2, g1)
                
                # 将 6 种 cost 堆叠起来，形状变为 [B, 6]
                costs = torch.stack([c0, c1, c2, c3, c4, c5], dim=1)
                # 找到每个 batch 中 cost 最小的那个排列的索引 (0 到 5)
                min_indices = torch.argmin(costs, dim=1)  # [B]
                
                # 预定义好 6 种排列对应的物理索引 (0代表g1, 1代表g2, 2代表g3)
                perms = torch.tensor([
                    [0, 1, 2],  # c0
                    [0, 2, 1],  # c1
                    [1, 0, 2],  # c2
                    [1, 2, 0],  # c3
                    [2, 0, 1],  # c4
                    [2, 1, 0]   # c5
                ], device=pred_xy.device)
                
                # 获取当前 batch 最优的排列索引组合，形状为 [B, 3]
                best_perms = perms[min_indices]
                
                # 将 Ground Truth 堆叠为 [B, 3, 2] 的统一张量
                g_all = torch.stack([g1, g2, g3], dim=1)
                
                # 利用高级索引一次性完成重组
                batch_idx = torch.arange(pred_xy.size(0), device=pred_xy.device).unsqueeze(-1)
                g_matched = g_all[batch_idx, best_perms]  # [B, 3, 2]
                
            # 2. 惩罚阶段：对重组后完美对齐的坐标计算 L2 Loss
            loss_p1 = loss_fn(p1, g_matched[:, 0]).sum(dim=-1)
            loss_p2 = loss_fn(p2, g_matched[:, 1]).sum(dim=-1)
            loss_p3 = loss_fn(p3, g_matched[:, 2]).sum(dim=-1)
            loss = (loss_p1 + loss_p2 + loss_p3).mean() / 3.0

            return loss
            
        return three_point_set_loss_smooth_l1

def get_mse(pred_xy, gt_xy):
    p1, p2, p3 = pred_xy[:,0], pred_xy[:,1], pred_xy[:,2]
    g1, g2, g3 = gt_xy[:,0], gt_xy[:,1], gt_xy[:,2]
    
    with torch.no_grad():
        def calc_cost(ga, gb, gc):
            return ((p1 - ga)**2).sum(dim=-1) + \
                   ((p2 - gb)**2).sum(dim=-1) + \
                   ((p3 - gc)**2).sum(dim=-1)
        c0 = calc_cost(g1, g2, g3)
        c1 = calc_cost(g1, g3, g2)
        c2 = calc_cost(g2, g1, g3)
        c3 = calc_cost(g2, g3, g1)
        c4 = calc_cost(g3, g1, g2)
        c5 = calc_cost(g3, g2, g1)
        costs = torch.stack([c0, c1, c2, c3, c4, c5], dim=1)
        min_costs, _ = torch.min(costs, dim=1)  # [B]
    mse_1d = min_costs.mean() / 6.0  # num_photons=3, num_axes=2
    return mse_1d

def train(model, trainLoader, optimizer, loss_fn):
    model.train()
    batchLoss = 0
    sum_squared_error_1d = 0; event_count = 0
    for batch_idx, (sample, label) in enumerate(trainLoader):
        sample, label = sample.cuda(), label.cuda()
        x1, y1, z1, e1 = label[:, 0, 0], label[:, 0, 1], label[:, 0, 2], label[:, 0, 3]
        x2, y2, z2, e2 = label[:, 1, 0], label[:, 1, 1], label[:, 1, 2], label[:, 1, 3]
        x3, y3, z3, e3 = label[:, 2, 0], label[:, 2, 1], label[:, 2, 2], label[:, 2, 3]
        gt_xy = torch.stack((torch.stack((x1, y1), axis=1), torch.stack((x2, y2), axis=1), torch.stack((x3, y3), axis=1)), axis=1)
        optimizer.zero_grad()
        output = model(sample)
        pred_xy = torch.stack((output[:,0:2], output[:,2:4], output[:,4:6]), axis=1)
        loss = loss_fn(pred_xy, gt_xy)
        loss.backward()
        optimizer.step()
        batchLoss += loss.item() * sample.shape[0]
        mse_1d = get_mse(pred_xy, gt_xy)
        sum_squared_error_1d += mse_1d.item() * sample.shape[0]
        event_count += sample.shape[0]
    avgLoss = batchLoss / len(trainLoader.dataset) / 6 ### divide by 6 to get the average loss per photon per axis
    rms_1d = np.sqrt(sum_squared_error_1d / event_count)
    print(f"[Train]\t Average Loss: {avgLoss:.6f} (RMS = {rms_1d:.6f})")
    return avgLoss, rms_1d

def evaluate(model, valLoader, loss_fn):
    model.eval()
    batchLoss = 0
    sum_squared_error_1d = 0; event_count = 0
    with torch.no_grad():
        for batch_idx, (sample, label) in enumerate(valLoader):
            sample, label = sample.cuda(), label.cuda()
            x1, y1, z1, e1 = label[:, 0, 0], label[:, 0, 1], label[:, 0, 2], label[:, 0, 3]
            x2, y2, z2, e2 = label[:, 1, 0], label[:, 1, 1], label[:, 1, 2], label[:, 1, 3]
            x3, y3, z3, e3 = label[:, 2, 0], label[:, 2, 1], label[:, 2, 2], label[:, 2, 3]
            gt_xy = torch.stack((torch.stack((x1, y1), axis=1), torch.stack((x2, y2), axis=1), torch.stack((x3, y3), axis=1)), axis=1)
            output = model(sample)
            pred_xy = torch.stack((output[:,0:2], output[:,2:4], output[:,4:6]), axis=1)
            loss = loss_fn(pred_xy, gt_xy)
            mse_1d = get_mse(pred_xy, gt_xy)
            sum_squared_error_1d += mse_1d.item() * sample.shape[0]
            event_count += sample.shape[0]
            batchLoss += loss.item() * sample.shape[0]
    avgLoss = batchLoss / len(valLoader.dataset) / 6 ### divide by 6 to get the average loss per photon per axis
    rms_1d = np.sqrt(sum_squared_error_1d / event_count)
    print(f"[Val]\t Average Loss: {avgLoss:.6f} (RMS = {rms_1d:.6f})")
    return avgLoss, rms_1d

def get_dataloaders(conf):
    """construct all dataloaders"""
    datasets = {}
    loaders = {}
    splits = ['Train', 'Val', 'Test']
    keys = ['train_files', 'val_files', 'test_files']
    batch_keys = ['batch_size_train', 'batch_size_val', 'batch_size_test']
    file_range_keys = ['train_file_range', 'val_file_range', 'test_file_range']

    for split, key, batch_key, file_range_key in zip(splits, keys, batch_keys, file_range_keys):
        files = [f"{conf.data.sample_folder}/pileupOf3phs_sample_{i}.npz" for i in range(conf.data[file_range_key][0], conf.data[file_range_key][1] + 1)]

        datasets[split] = triplePhotonDataset(
            files,
            sampleRatio = conf.data.sample_ratio,
            datasetName = split.capitalize(),
        )

        loaders[split] = torch.utils.data.DataLoader(
            datasets[split],
            batch_size=conf.data[batch_key],
            shuffle=(split=='Train'),
            num_workers=conf.data.num_workers,
            pin_memory=True
        )
    return loaders['Train'], loaders['Val'], loaders['Test']

def plot_loss_curves(train_losses, val_losses, test_loss, exp_name, conf):
    import matplotlib.pyplot as plt
    plt.figure(figsize=(8,6))
    plt.plot(train_losses, label='Train Loss')
    plt.plot(val_losses, label='Val Loss')
    if test_loss > 0:
        plt.axhline(y=test_loss, color='green', linestyle='--', label='Test Loss')
    plt.xlabel('Epoch')
    plt.ylabel('MSE Loss')
    plt.yscale('log')
    plt.legend()
    plt.grid()
    plotName = f'loss_curve_triplePhoton_{conf.model.version}.png'
    plt.savefig(f'Results/3ph/{exp_name}/Plots/{plotName}')
    plt.close()

def plot_rms_curves(train_rms, val_rms, test_rms, exp_name, conf):
    import matplotlib.pyplot as plt
    plt.figure(figsize=(8,6))
    plt.plot(train_rms, label='Train RMS')
    plt.plot(val_rms, label='Val RMS')
    if test_rms > 0:
        plt.axhline(y=test_rms, color='green', linestyle='--', label='Test RMS')
    plt.xlabel('Epoch')
    plt.ylabel('RMS Error [pixels]')
    plt.yscale('log')
    plt.legend()
    plt.grid()
    plotName = f'rms_curve_triplePhoton_{conf.model.version}.png'
    plt.savefig(f'Results/3ph/{exp_name}/Plots/{plotName}')
    plt.close()

def get_model_name(conf):
    modelName = f'triplePhoton{conf.model.version}_{conf.data.energy}keV'
    return modelName

if __name__ == "__main__":
    exp_name = prepare_output_folder(conf)
    model = get_triple_photon_model_class(conf.model.version)().cuda()
    # summary(model, input_size=(128, 3, conf.data.n_size, conf.data.n_size))
    
    loss_fn = get_loss_function(conf)
    optimizer = torch.optim.Adam(model.parameters(), lr=conf.training.learning_rate, weight_decay=conf.training.weight_decay)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=conf.training.scheduler_factor, patience=conf.training.scheduler_patience)

    trainLoader, valLoader, testLoader = get_dataloaders(conf)

    TrainLosses, ValLosses = [], []
    train_rms, val_rms = [], []

    for epoch in tqdm(range(1, conf.training.epochs + 1)):
        train_loss, train_rms_epoch = train(model, trainLoader, optimizer, loss_fn)
        val_loss, val_rms_epoch = evaluate(model, valLoader, loss_fn)
        TrainLosses.append(train_loss)
        ValLosses.append(val_loss)
        train_rms.append(train_rms_epoch)
        val_rms.append(val_rms_epoch)
        scheduler.step(val_loss)
        print(f"Learning Rate: {optimizer.param_groups[0]['lr']:.2e}")
        if epoch in conf.training.checkpoint_epochs or epoch == conf.training.epochs:
            modelName = get_model_name(conf)
            torch.save(model.state_dict(), f'Results/3ph/{exp_name}/Models/{modelName}_E{epoch}.pth')
            print(f"Saved model checkpoint: {modelName}_E{epoch}.pth")
            plot_loss_curves(TrainLosses, ValLosses, test_loss=-1, exp_name=exp_name, conf=conf)
            plot_rms_curves(train_rms, val_rms, test_rms=-1, exp_name=exp_name, conf=conf)
    test_loss, test_rms = evaluate(model, testLoader, loss_fn)
    plot_loss_curves(TrainLosses, ValLosses, test_loss=test_loss, exp_name=exp_name, conf=conf)
    plot_rms_curves(train_rms, val_rms, test_rms=test_rms, exp_name=exp_name, conf=conf)