import torch
import sys
sys.path.append('./src')
import models
from datasets import *
from tqdm import tqdm
from matplotlib import pyplot as plt

torch.manual_seed(42)
torch.cuda.manual_seed(42)
np.random.seed(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

configs = {}
configs['SiemenStar'] = {
    'dataFiles': [f'/mnt/sls_det_storage/moench_data/MLXID/Samples/Measurement/2504_SOLEIL_SiemenStarClusters_MOENCH040_150V/clusters_chunk{i}.h5' for i in range(200)],
    'modelVersion': '251022',
    'roi': [140, 230, 120, 210], # x_min, x_max, y_min, y_max,
    'noise': 0.13  # in keV
}
BinningFactor = 10
numberOfAugOps = 6
Roi = configs['SiemenStar']['roi']
X_st, X_ed, Y_st, Y_ed = Roi
mlSuperFrame = np.zeros(((Y_ed-Y_st)*BinningFactor, (X_ed-X_st)*BinningFactor))
countFrame = np.zeros((Y_ed-Y_st, X_ed-X_st))
subpixelDistribution = np.zeros((BinningFactor, BinningFactor))

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)

if __name__ == "__main__":
    task = 'SiemenStar'
    config = configs[task]

    model = models.get_model_class(config['modelVersion'])().cuda()
    model.load_state_dict(torch.load(f'/home/xie_x1/MLXID/DeepLearning/Models/singlePhoton{config["modelVersion"]}_15.3keV_Noise{config["noise"]}keV_E500_aug8.pth', weights_only=True))
    predictions = []
    referencePoints = []
    nChunks = len(config['dataFiles']) // 32 + 1
    for idxChunk in range(nChunks):
        stFileIdx = idxChunk * 32
        edFileIdx = min((idxChunk + 1) * 32, len(config['dataFiles']))
        sampleFiles = config['dataFiles'][stFileIdx : edFileIdx]
        print(f'Processing files {stFileIdx} to {edFileIdx}...')
        dataset = singlePhotonDataset(
            sampleFiles, 
            sampleRatio=1.0, 
            datasetName='Inference',
            numberOfAugOps=numberOfAugOps
            )
        dataLoader = torch.utils.data.DataLoader(
            dataset,
            batch_size=8192,
            shuffle=False,
            num_workers=32,
            pin_memory=True,
        )

        referencePoints.append(dataset.referencePoint)

        with torch.no_grad():
            for batch in tqdm(dataLoader):
                inputs, _ = batch
                inputs = inputs.cuda()
                outputs = model(inputs)[:, :2]  # only x and y
                predictions.append(outputs.cpu())

    predictions = torch.cat(predictions, dim=0)
    predictions = apply_inverse_transforms(predictions, numberOfAugOps)
    predictions += torch.tensor([1.5, 1.5]).unsqueeze(0)  # adjust back to original coordinate system
    referencePoints = np.concatenate(referencePoints, axis=0)
    print(f'mean x = {torch.mean(predictions[:, 0])}, std x = {torch.std(predictions[:, 0])}')
    print(f'mean y = {torch.mean(predictions[:, 1])}, std y = {torch.std(predictions[:, 1])}')
    absolutePositions = predictions.numpy() + referencePoints[:, :2] - 1
    
    hit_x = np.floor((absolutePositions[:, 0] - Roi[0]) * BinningFactor).astype(int)
    hit_x = np.clip(hit_x, 0, mlSuperFrame.shape[1]-1)
    hit_y = np.floor((absolutePositions[:, 1] - Roi[2]) * BinningFactor).astype(int)
    hit_y = np.clip(hit_y, 0, mlSuperFrame.shape[0]-1)
    np.add.at(mlSuperFrame, (hit_y, hit_x), 1)

    np.add.at(countFrame, ((referencePoints[:, 1] - Roi[2]).astype(int),
                           (referencePoints[:, 0] - Roi[0]).astype(int)), 1)

    np.add.at(subpixelDistribution,
              (np.floor((absolutePositions[:, 1] % 1) * BinningFactor).astype(int),
               np.floor((absolutePositions[:, 0] % 1) * BinningFactor).astype(int)), 1)

    plt.imshow(mlSuperFrame, origin='lower', extent=[Y_st, Y_ed, X_st, X_ed])
    plt.colorbar()
    plt.savefig('InferenceResults/SiemenStar_ML_superFrame.png', dpi=300)
    np.save('InferenceResults/SiemenStar_ML_superFrame.npy', mlSuperFrame)
    plt.clf()

    plt.imshow(countFrame, origin='lower', extent=[Y_st, Y_ed, X_st, X_ed])
    plt.colorbar()
    plt.savefig('InferenceResults/SiemenStar_count_Frame.png', dpi=300)
    np.save('InferenceResults/SiemenStar_count_Frame.npy', countFrame)

    plt.clf()
    plt.imshow(subpixelDistribution, origin='lower')
    plt.colorbar()
    plt.savefig('InferenceResults/SiemenStar_subpixel_Distribution.png', dpi=300)
    np.save('InferenceResults/SiemenStar_subpixel_Distribution.npy', subpixelDistribution)