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

from etaInterpolationFromPoints import interpolate_eta_from_points

from pathlib import Path
from omegaconf import OmegaConf
import torch
from tqdm import tqdm
from matplotlib import pyplot as plt
import numpy as np
import h5py

from models import get_model_class
from datasets import singlePhotonDataset

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/infer_1photon.yaml")
NX, NY = conf.data.names[conf.experiment.name].NX, conf.data.names[conf.experiment.name].NY

def inv0(p): return p
def inv1(p): return torch.stack([-p[..., 0], p[..., 1]], dim=-1)
def inv2(p): return torch.stack([p[..., 0], -p[..., 1]], dim=-1)
def inv3(p): return -p
def inv4(p): return torch.stack([p[..., 1], p[..., 0]], dim=-1)
def inv5(p): return torch.stack([p[..., 1], -p[..., 0]], dim=-1) 
def inv6(p): return torch.stack([-p[..., 1], p[..., 0]], dim=-1)
def inv7(p): return torch.stack([-p[..., 1], -p[..., 0]], dim=-1)

INVERSE_TRANSFORMS = {
    0: inv0, 1: inv1, 2: inv2, 3: inv3, 
    4: inv4, 5: inv5, 6: inv6, 7: inv7,
}

def apply_inverse_transforms(predictions: torch.Tensor, numberOfAugOps: int) -> torch.Tensor:
    N = predictions.shape[0] // numberOfAugOps
    preds = predictions.view(N, numberOfAugOps, 2)
    corrected = torch.zeros_like(preds)
    for idx in range(numberOfAugOps):
        corrected[:, idx, :] = INVERSE_TRANSFORMS[idx](preds[:, idx, :])
    return corrected.mean(dim=1)

def prepare_output_folder(conf):
    output_base = Path(conf.experiment.output_base) / conf.experiment.name / (f'1ph_{conf.model.experiment_name}') / (f'augX{conf.inference.num_aug_ops}')
    output_base.mkdir(parents=True, exist_ok=True)
    OmegaConf.save(conf, output_base / 'config.yaml')
    return output_base

def get_files_list(conf):
    task_conf = conf.data.names[conf.experiment.name]
    file_pattern = task_conf.file_pattern
    start, end = task_conf.file_range
    files = [str(Path(conf.data.sample_folder) / file_pattern.format(i)) for i in range(start, end)]
    return files

def run_inference(model, data_loader, conf):
    all_predictions = []
    all_reference_points = data_loader.dataset.referencePoint
    with torch.no_grad():
        for batch in tqdm(data_loader, desc="Inferring"):
            inputs, _ = batch
            inputs = inputs.to('cuda')
            outputs = model(inputs)[:, :2].cpu()  # 只取 x, y
            all_predictions.append(outputs)

    all_predictions = torch.cat(all_predictions, dim=0)
    all_predictions = apply_inverse_transforms(all_predictions, conf.inference.num_aug_ops)
    offset = [inputs.shape[-2] / 2., inputs.shape[-1] / 2.]
    offset = torch.tensor(offset).unsqueeze(0)  # (1, 2)
    all_predictions += offset
    print(f'[Inference]: mean x = {torch.mean(all_predictions[:, 0]):.4f}, std x = {torch.std(all_predictions[:, 0]):.4f}')
    print(f'[Inference]: mean y = {torch.mean(all_predictions[:, 1]):.4f}, std y = {torch.std(all_predictions[:, 1]):.4f}')
    return all_predictions.numpy(), all_reference_points

def accumulate_hits(predictions: np.ndarray, reference_points: np.ndarray, 
                    binning_factor: int, number_of_subframes: int = 20):
    ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor), dtype=np.uint32)
    count_frame = np.zeros((NY, NX), dtype=np.uint32)
    subpixel_dist = np.zeros((binning_factor, binning_factor), dtype=np.uint32)
    
    chunk_size = len(predictions) // number_of_subframes
    
    for i in range(number_of_subframes):
        start_idx = i * chunk_size
        end_idx = (i + 1) * chunk_size if i < number_of_subframes - 1 else len(predictions)
        
        pred_chunk = predictions[start_idx:end_idx]
        ref_chunk = reference_points[start_idx:end_idx]
        
        abs_pos_chunk = pred_chunk + ref_chunk[:, :2]
        
        # --- Super resolution frames ---
        hit_x = np.floor(abs_pos_chunk[:, 0] * binning_factor).astype(np.int32)
        hit_y = np.floor(abs_pos_chunk[:, 1] * binning_factor).astype(np.int32)
        np.add.at(ml_super_frames[i], (hit_y, hit_x), 1)
        
        # --- Count frame ---
        ref_x = (ref_chunk[:, 0] + 1).astype(np.int32)
        ref_y = (ref_chunk[:, 1] + 1).astype(np.int32)
        np.add.at(count_frame, (ref_y, ref_x), 1)
        
        # --- Subpixel distribution ---
        sub_x = np.floor((abs_pos_chunk[:, 0] % 1) * binning_factor).astype(np.int32)
        sub_y = np.floor((abs_pos_chunk[:, 1] % 1) * binning_factor).astype(np.int32)
        np.add.at(subpixel_dist, (sub_y, sub_x), 1)
        
    return ml_super_frames, count_frame, subpixel_dist

def save_results(ml_super_frames, ml_super_frames_eta,
                count_frame, subpixel_dist, subpixel_dist_eta,
                roi: list, binning_factor: int, output_dir: Path):
    x_min, x_max, y_min, y_max = roi
    
    # 1. super-resolution frame
    np.save(output_dir / '1Photon_ML_superFrames.npy', ml_super_frames)

    ml_super_frame = np.sum(ml_super_frames, axis=0)
    plt.figure(figsize=(8, 8))
    plt.imshow(ml_super_frame[y_min*binning_factor:y_max*binning_factor, x_min*binning_factor:x_max*binning_factor], origin='lower', extent=[x_min, x_max, y_min, y_max])
    plt.colorbar(label='Counts')
    plt.title('ML Super-Resolution Frame')
    plt.xlabel('X (pixel)')
    plt.ylabel('Y (pixel)')
    plt.savefig(output_dir / '1Photon_ML_superFrame.png', dpi=300, bbox_inches='tight')
    plt.close()
    np.save(output_dir / '1Photon_ML_superFrame.npy', ml_super_frame)

    # 2. super-resolution frame with eta interpolation
    np.save(output_dir / '1Photon_ML_superFrames_etaInterpolated.npy', ml_super_frames_eta)

    ml_super_frame_eta = np.sum(ml_super_frames_eta, axis=0)
    plt.figure(figsize=(8, 8))
    plt.imshow(ml_super_frame_eta[y_min*binning_factor:y_max*binning_factor, x_min*binning_factor:x_max*binning_factor], origin='lower', extent=[x_min, x_max, y_min, y_max])
    plt.colorbar(label='Counts')
    plt.title('ML Super-Resolution Frame with Eta Interpolation')
    plt.xlabel('X (pixel)')
    plt.ylabel('Y (pixel)')
    plt.savefig(output_dir / '1Photon_ML_superFrame_etaInterpolated.png', dpi=300, bbox_inches='tight')
    plt.close()
    np.save(output_dir / '1Photon_ML_superFrame_etaInterpolated.npy', ml_super_frame_eta)

    # 3. count frame
    plt.figure(figsize=(8, 8))
    plt.imshow(count_frame[y_min:y_max, x_min:x_max], origin='lower', extent=[x_min, x_max, y_min, y_max])
    plt.colorbar(label='Counts')
    plt.title('Photon Count Frame')
    plt.xlabel('X (pixel)')
    plt.ylabel('Y (pixel)')
    plt.savefig(output_dir / '1Photon_count_Frame.png', dpi=300, bbox_inches='tight')
    plt.close()
    np.save(output_dir / '1Photon_count_Frame.npy', count_frame)
    
    # 4. sub-pixel distribution
    plt.figure(figsize=(8, 8))
    plt.imshow(subpixel_dist, origin='lower', 
              extent=[0, binning_factor, 0, binning_factor],
              cmap='viridis')
    plt.colorbar(label='Counts')
    plt.title('Sub-pixel Distribution')
    plt.xlabel(f'Sub-pixel X (1/{binning_factor} pixel)')
    plt.ylabel(f'Sub-pixel Y (1/{binning_factor} pixel)')
    plt.savefig(output_dir / '1Photon_subpixel_Distribution.png', dpi=300, bbox_inches='tight')
    plt.close()
    np.save(output_dir / '1Photon_subpixel_Distribution.npy', subpixel_dist)
    rms, mean = np.std(subpixel_dist), np.mean(subpixel_dist)
    print(f"[Plotting]: Sub-pixel distribution: RMS/Mean: {rms/mean:.4f}, expected value = {1/np.sqrt(mean):.4f} for uniform distribution")

    # 5. sub-pixel distribution with eta interpolation
    plt.figure(figsize=(8, 8))
    plt.imshow(subpixel_dist_eta, origin='lower',
              extent=[0, binning_factor, 0, binning_factor],
              cmap='viridis')
    plt.colorbar(label='Counts')
    plt.title('Sub-pixel Distribution with Eta Interpolation')
    plt.xlabel(f'Sub-pixel X (1/{binning_factor} pixel)')
    plt.ylabel(f'Sub-pixel Y (1/{binning_factor} pixel)')
    plt.savefig(output_dir / '1Photon_subpixel_Distribution_etaInterpolated.png', dpi=300, bbox_inches='tight')
    plt.close()
    np.save(output_dir / '1Photon_subpixel_Distribution_etaInterpolated.npy', subpixel_dist_eta)
    rms_eta, mean_eta = np.std(subpixel_dist_eta), np.mean(subpixel_dist_eta)
    print(f"[Plotting]: Sub-pixel distribution with eta interpolation: RMS/Mean: {rms_eta/mean_eta:.4f}, expected value = {1/np.sqrt(mean_eta):.4f} for uniform distribution")
    
    print(f"[Plotting]: Results saved to: {output_dir}")

if __name__ == "__main__":
    ### output preparation
    output_dir = prepare_output_folder(conf)

    ### model loading
    model_version = conf.model.experiment_name.split('_v')[-1][:6]
    model = get_model_class(model_version)().cuda()
    model.load_state_dict(torch.load(f'{conf.model.base_dir}/{conf.model.experiment_name}/Models/{conf.model.name}', weights_only=True))
    model.eval()

    ### data loading
    files_list = get_files_list(conf)
    roi = conf.data.names[conf.experiment.name].roi
    BinningFactor = conf.inference.binning_factor
    numberOfAugOps = conf.inference.num_aug_ops
    flag_normalize = conf.data.normalize
    nChunks = int(np.ceil(len(files_list) / conf.inference.chunk_size))

    list_of_predictions = []
    list_of_reference_points = []

    ml_super_frame = np.zeros((NY * BinningFactor, NX * BinningFactor))
    count_frame = np.zeros((NY, NX))
    subpixel_dist = np.zeros((BinningFactor, BinningFactor))

    ### Inference loop
    for idxChunk in range(nChunks):
        start_idx = idxChunk * conf.inference.chunk_size
        end_idx = min(start_idx + conf.inference.chunk_size, len(files_list))
        chunk_files = files_list[start_idx:end_idx]
        print(f'[Inferring] Chunk {idxChunk+1}/{nChunks}: Loading files {start_idx} to {end_idx}...')

        dataset = singlePhotonDataset(
            chunk_files,
            sampleRatio=1.0,
            datasetName='Inference',
            numberOfAugOps=conf.inference.num_aug_ops,
            normalize=conf.data.normalize
        )
        
        data_loader = torch.utils.data.DataLoader(
            dataset,
            batch_size=conf.data.batch_size,
            shuffle=False,
            num_workers=conf.data.num_workers,
            pin_memory=True
        )
        
        _predictions, _ref_points = run_inference(model, data_loader, conf)
        list_of_predictions.append(_predictions)
        list_of_reference_points.append(_ref_points)

        del dataset, data_loader
        torch.cuda.empty_cache()

    predictions = np.concatenate(list_of_predictions, axis=0)
    ref_points = np.concatenate(list_of_reference_points, axis=0)

    del list_of_predictions, list_of_reference_points

    ml_super_frames, count_frame, subpixel_dist = accumulate_hits(
        predictions, ref_points, binning_factor=BinningFactor
    )

    print('[Main]: Applying eta interpolation to predictions...')
    predictions_eta_interpolated = interpolate_eta_from_points(predictions)
    ml_super_frames_eta, count_frame_eta, subpixel_dist_eta = accumulate_hits(
        predictions_eta_interpolated, ref_points, binning_factor=BinningFactor
    )

    save_results(ml_super_frames, ml_super_frames_eta,
                    count_frame, subpixel_dist, subpixel_dist_eta,
                    roi=roi, binning_factor=BinningFactor, 
                    output_dir=output_dir
                )