# Liam McLauglin
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from scipy.fftpack import idctn
import cv2
from sklearn.metrics import confusion_matrix
import timm
import warnings
warnings.filterwarnings('ignore')

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {DEVICE}")

KEEP_LABELS    = {0, 2, 3, 5, 6, 8}
LABEL_REMAP    = {0: 0, 2: 1, 3: 2, 5: 3, 6: 4, 8: 5}
LABEL_UNMAP    = {v: k for k, v in LABEL_REMAP.items()}
SCORED_CLASSES = [0, 1, 2, 3, 4]
BCKG_CLASS     = 5
N_CLASSES      = 6

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)
WEIGHTS_PATH  = "../data/vit_b16.pth"
DATA_DIR      = "../data"
OUT_DIR       = "../data"


# Computes the weighted competition score from per-class error rates on scored classes and the background class.
def compute_score(y_true, y_pred):
    cm = confusion_matrix(y_true, y_pred, labels=list(range(N_CLASSES)))
    errors = []
    for c in SCORED_CLASSES:
        row = cm[c]; total = row.sum()
        err = (total - cm[c, c]) / total if total > 0 else 0.0
        errors.append(err)
    avg_lbl  = np.mean(errors)
    bt       = cm[BCKG_CLASS].sum()
    bckg_err = (bt - cm[BCKG_CLASS, BCKG_CLASS]) / bt if bt > 0 else 0.0
    score    = 0.9 * avg_lbl + 0.1 * bckg_err
    return avg_lbl * 100, bckg_err * 100, score * 100


# Loads a CSV file and filters rows to the kept labels, remapping them to contiguous class indices.
def load_and_filter(path, is_eval=False):
    print(f"Loading {path} ...")
    df = pd.read_csv(path, comment='#', header=None)
    label_col    = 0
    feature_cols = list(range(1, df.shape[1]))
    if is_eval:
        X = df[feature_cols].values.astype(np.float32)
        return X, np.zeros(len(df), dtype=int), np.arange(len(df)), len(df)
    mask = df[label_col].isin(KEEP_LABELS)
    df_f = df[mask].reset_index(drop=True)
    idx  = np.where(mask.values)[0]
    y    = np.array([LABEL_REMAP[l] for l in df_f[label_col].values], dtype=int)
    X    = df_f[feature_cols].values.astype(np.float32)
    print(f"  {len(X)} samples | {dict(zip(*np.unique(y, return_counts=True)))}")
    return X, y, idx, len(df)


# Converts a row of DCT coefficients to a normalized spatial RGB image via inverse DCT and ImageNet normalization.
def row_to_spatial_image(row, size=224):
    channels = []
    for ch in range(3):
        coeff   = row[ch * 1024: (ch + 1) * 1024].reshape(32, 32)
        spatial = idctn(coeff, norm='ortho')
        channels.append(spatial)
    img = np.stack(channels, axis=-1).astype(np.float32)
    mn, mx = img.min(), img.max()
    img = (img - mn) / (mx - mn) * 255.0 if mx > mn else np.zeros_like(img)
    img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
    img = (img / 255.0 - IMAGENET_MEAN) / IMAGENET_STD
    return torch.tensor(img.transpose(2, 0, 1), dtype=torch.float32)


# Reshapes a row of raw DCT coefficients directly into a normalized image without inverse DCT.
def row_to_dct_image(row, size=224):
    img = row.reshape(3, 32, 32).transpose(1, 2, 0).astype(np.float32)
    mn, mx = img.min(), img.max()
    img = (img - mn) / (mx - mn) * 255.0 if mx > mn else np.zeros_like(img)
    img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
    img = (img / 255.0 - IMAGENET_MEAN) / IMAGENET_STD
    return torch.tensor(img.transpose(2, 0, 1), dtype=torch.float32)


# Dataset wrapper that converts feature rows into image tensors using either spatial or DCT mode.
class PatchDataset(Dataset):
    # Stores the feature matrix, optional labels, and the image conversion mode.
    def __init__(self, X, y=None, mode='spatial'):
        self.X    = X
        self.y    = y
        self.mode = mode
    # Returns the number of samples in the dataset.
    def __len__(self):
        return len(self.X)
    # Returns the converted image tensor (and label if available) for the given index.
    def __getitem__(self, idx):
        img = row_to_spatial_image(self.X[idx]) if self.mode == 'spatial' \
              else row_to_dct_image(self.X[idx])
        if self.y is not None:
            return img, self.y[idx]
        return img


# Computes inverse-frequency class weights to balance the loss across imbalanced classes.
def inverse_class_weights(y):
    counts = np.bincount(y, minlength=N_CLASSES)
    return np.where(counts > 0, len(y) / (N_CLASSES * counts), 1.0)


# Builds a PyTorch DataLoader from features and labels with the given batch size, shuffle, and image mode.
def get_loader(X, y=None, batch_size=16, shuffle=False, mode='spatial'):
    ds = PatchDataset(X, y, mode=mode)
    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle,
                      num_workers=4, pin_memory=True)


# Constructs a ViT-Base/16 model, loads pretrained weights, and replaces the head with a fresh classifier for the target classes.
def build_vit(drop_path=0.1):
    print(f"  Loading ViT weights from {WEIGHTS_PATH}...")
    model = timm.create_model('vit_base_patch16_224',
                               pretrained=False,
                               drop_path_rate=drop_path)
    state = torch.load(WEIGHTS_PATH, map_location='cpu')
    missing, unexpected = model.load_state_dict(state, strict=False)
    print(f"  Weights loaded — missing: {len(missing)}, unexpected: {len(unexpected)}")
    in_feat    = model.head.in_features
    model.head = nn.Linear(in_feat, N_CLASSES)
    return model.to(DEVICE)


# Builds optimizer parameter groups with layer-wise learning-rate decay and selective weight decay (no decay on biases/norms).
def create_param_groups_lrd(model, lr, layer_decay=0.75, weight_decay=0.05):
    param_groups = []
    try:
        num_layers = len(model.blocks) + 1
    except AttributeError:
        num_layers = 12
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if 'head' in name or 'norm' in name:
            layer_id = num_layers
        elif 'patch_embed' in name:
            layer_id = 0
        elif 'blocks' in name:
            try:
                block_num = int(name.split('blocks.')[1].split('.')[0])
                layer_id  = block_num + 1
            except (IndexError, ValueError):
                layer_id = num_layers // 2
        else:
            layer_id = 0
        lr_scale = layer_decay ** (num_layers - layer_id)
        wd = 0.0 if ('bias' in name or 'norm' in name) else weight_decay
        param_groups.append({'params': [param],
                             'lr': lr * lr_scale,
                             'weight_decay': wd})
    return param_groups


# Runs the model in eval mode over a loader and returns stacked softmax probabilities for every sample.
def get_probs(model, loader):
    model.eval()
    all_probs = []
    with torch.no_grad():
        for batch in loader:
            imgs  = batch[0] if isinstance(batch, (list, tuple)) else batch
            probs = torch.softmax(model(imgs.to(DEVICE)), dim=1)
            all_probs.append(probs.cpu().numpy())
    return np.vstack(all_probs)


# Trains a ViT on the training set with class-balanced loss, OneCycle LR, and smoothed early stopping on the early-stop split.
def train_vit(X_tr, y_tr, X_es, y_es, mode,
              lr=1.98e-4, weight_decay=0.0203,
              layer_decay=0.859, drop_path=0.190,
              max_epochs=20, patience=5, min_epochs=8,
              smooth_window=3, batch_size=16, model_name="ViT"):
    print(f"\n[{model_name}] Training on train only (mode={mode})...")
    cw_t  = torch.tensor(inverse_class_weights(y_tr),
                         dtype=torch.float32).to(DEVICE)
    tr_ld = get_loader(X_tr, y_tr, batch_size, shuffle=True,  mode=mode)
    es_ld = get_loader(X_es, y_es, batch_size, shuffle=False, mode=mode)
    model     = build_vit(drop_path)
    optimizer = torch.optim.AdamW(
        create_param_groups_lrd(model, lr, layer_decay, weight_decay))
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=lr,
        epochs=max_epochs,
        steps_per_epoch=len(tr_ld),
        pct_start=0.1)
    criterion  = nn.CrossEntropyLoss(weight=cw_t)
    best_score, best_state, best_epoch = float('inf'), None, 0
    recent_scores = []
    no_improve    = 0

    for epoch in range(max_epochs):
        model.train()
        total_loss = 0.0
        for imgs, labels in tr_ld:
            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
            optimizer.zero_grad()
            loss = criterion(model(imgs), labels)
            loss.backward()
            optimizer.step()
            scheduler.step()
            total_loss += loss.item()

        probs = get_probs(model, es_ld)
        _, _, score = compute_score(y_es, np.argmax(probs, axis=1))
        recent_scores.append(score)
        smoothed = np.mean(recent_scores[-smooth_window:])
        print(f"  Epoch {epoch+1:02d} | loss={total_loss/len(tr_ld):.4f} "
              f"| es={score:.4f}% | smooth={smoothed:.4f}%")

        if smoothed < best_score and epoch + 1 >= smooth_window:
            best_score = smoothed
            best_state = {k: v.clone() for k, v in model.state_dict().items()}
            best_epoch = epoch + 1
            no_improve = 0
        elif epoch + 1 >= min_epochs:
            no_improve += 1
            if no_improve >= patience:
                print(f"  Early stopping at epoch {epoch+1}")
                break

    model.load_state_dict(best_state)
    print(f"  [{model_name}] Best smoothed dev: {best_score:.4f}% at epoch {best_epoch}")
    return model, best_score, best_epoch


# Retrains a fresh ViT on combined train+dev data for a fixed number of epochs (no early stopping).
def retrain_vit(X_all, y_all, mode, best_epochs,
                lr=1.98e-4, weight_decay=0.0203,
                layer_decay=0.859, drop_path=0.190,
                batch_size=16, model_name="ViT"):
    print(f"\n[{model_name}] Retraining on train+dev ({best_epochs} epochs)...")
    cw_t  = torch.tensor(inverse_class_weights(y_all),
                         dtype=torch.float32).to(DEVICE)
    ld    = get_loader(X_all, y_all, batch_size, shuffle=True, mode=mode)
    model     = build_vit(drop_path)
    optimizer = torch.optim.AdamW(
        create_param_groups_lrd(model, lr, layer_decay, weight_decay))
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=lr,
        epochs=best_epochs,
        steps_per_epoch=len(ld),
        pct_start=0.1)
    criterion = nn.CrossEntropyLoss(weight=cw_t)
    for epoch in range(best_epochs):
        model.train()
        total_loss = 0.0
        for imgs, labels in ld:
            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
            optimizer.zero_grad()
            loss = criterion(model(imgs), labels)
            loss.backward()
            optimizer.step()
            scheduler.step()
            total_loss += loss.item()
        if (epoch + 1) % 5 == 0:
            print(f"  Epoch {epoch+1:02d} | loss={total_loss/len(ld):.4f}")
    return model


# Grid-searches the temperature value that minimizes the competition score after softmax rescaling.
def find_temperature(probs, y_true):
    print("\n[Post] Temperature scaling...")
    best_T, best_s = 1.0, float('inf')
    for T in np.arange(0.1, 3.0, 0.05):
        lp  = np.log(probs + 1e-10) / T
        ep  = np.exp(lp - lp.max(axis=1, keepdims=True))
        cal = ep / ep.sum(axis=1, keepdims=True)
        _, _, s = compute_score(y_true, np.argmax(cal, axis=1))
        if s < best_s:
            best_s, best_T = s, T
    print(f"  Best T={best_T:.2f} -> {best_s:.4f}%")
    return best_T


# Applies temperature scaling to a probability matrix and returns renormalized probabilities.
def apply_temperature(probs, T):
    lp = np.log(probs + 1e-10) / T
    ep = np.exp(lp - lp.max(axis=1, keepdims=True))
    return ep / ep.sum(axis=1, keepdims=True)


# Greedy per-class threshold search that divides probabilities by class thresholds to minimize the score.
def tune_thresholds(probs, y_true):
    print("\n[Post] Per-class threshold tuning...")
    best_thr, best_s = np.ones(N_CLASSES), float('inf')
    for c in range(N_CLASSES):
        for t in np.arange(0.3, 2.0, 0.05):
            thr   = best_thr.copy(); thr[c] = t
            preds = np.argmax(probs / thr[np.newaxis, :], axis=1)
            _, _, s = compute_score(y_true, preds)
            if s < best_s:
                best_s = s; best_thr = thr.copy()
    print(f"  Thresholds: {np.round(best_thr, 3)} -> {best_s:.4f}%")
    return best_thr


# Produces final class predictions by argmax over threshold-rescaled probabilities.
def predict_final(probs, thr):
    return np.argmax(probs / thr[np.newaxis, :], axis=1)


# Maps remapped predictions back to original label space and writes the full-length submission CSV.
def write_submission(preds, keep_idx, total_len, path):
    output = np.zeros(total_len, dtype=int)
    for i, p in zip(keep_idx, preds):
        output[i] = LABEL_UNMAP[p]
    with open(path, 'w') as f:
        f.write("label\n")
        for label in output:
            f.write(f"{label}\n")
    print(f"  Saved: {path}")


# Orchestrates the full pipeline: load data, train two ViTs, blend, calibrate, retrain on all data, and write submissions.
def main():
    TRAIN_PATH = f"{DATA_DIR}/train.csv"
    DEV_PATH   = f"{DATA_DIR}/dev.csv"
    EVAL_PATH  = f"{DATA_DIR}/eval.csv"

    X_tr,   y_tr,   tr_idx,   tr_len   = load_and_filter(TRAIN_PATH)
    X_dev,  y_dev,  dev_idx,  dev_len  = load_and_filter(DEV_PATH)
    X_eval, y_eval, eval_idx, eval_len = load_and_filter(EVAL_PATH, is_eval=True)

    print(f"\nTrain: {len(X_tr)}  Dev: {len(X_dev)}  Eval: {len(X_eval)}")

    rng      = np.random.default_rng(42)
    dev_perm = rng.permutation(len(y_dev))
    n_es     = int(0.7 * len(y_dev))
    es_idx   = dev_perm[:n_es]
    cal_idx  = dev_perm[n_es:]

    X_es,  y_es  = X_dev[es_idx],  y_dev[es_idx]
    X_cal, y_cal = X_dev[cal_idx], y_dev[cal_idx]
    print(f"  Dev split — ES: {len(y_es)}, Cal: {len(y_cal)}")

    X_all = np.vstack([X_tr, X_dev])
    y_all = np.concatenate([y_tr, y_dev])

    vit_kwargs = dict(
        lr=1.98e-4, weight_decay=0.0203,
        layer_decay=0.859, drop_path=0.190,
        max_epochs=20, patience=5,
        min_epochs=8, smooth_window=3,
        batch_size=16
    )

    model_a, score_a, epochs_a = train_vit(
        X_tr, y_tr, X_es, y_es,
        mode='spatial', model_name="ViT-A (spatial)", **vit_kwargs
    )
    model_b, score_b, epochs_b = train_vit(
        X_tr, y_tr, X_es, y_es,
        mode='dct', model_name="ViT-B (dct)", **vit_kwargs
    )

    print("\n[Blend] Evaluating on cal split...")
    cal_a = get_loader(X_cal, y_cal, 16, mode='spatial')
    cal_b = get_loader(X_cal, y_cal, 16, mode='dct')
    pa    = get_probs(model_a, cal_a)
    pb    = get_probs(model_b, cal_b)
    _, _, s_a = compute_score(y_cal, np.argmax(pa, axis=1))
    _, _, s_b = compute_score(y_cal, np.argmax(pb, axis=1))
    print(f"  ViT-A (spatial): {s_a:.4f}%")
    print(f"  ViT-B (dct):     {s_b:.4f}%")

    if s_a > 2 * s_b:
        print("  ViT-A too weak — dropping from blend")
        dev_blend = pb
        w_a, w_b  = 0.0, 1.0
    elif s_b > 2 * s_a:
        print("  ViT-B too weak — dropping from blend")
        dev_blend = pa
        w_a, w_b  = 1.0, 0.0
    else:
        w_a = (1.0/s_a) / (1.0/s_a + 1.0/s_b)
        w_b = 1.0 - w_a
        dev_blend = w_a * pa + w_b * pb

    print(f"  Weights: A={w_a:.3f} B={w_b:.3f}")
    _, _, s_blend = compute_score(y_cal, np.argmax(dev_blend, axis=1))
    print(f"  Blended: {s_blend:.4f}%")

    best_T   = find_temperature(dev_blend, y_cal)
    cal_cal  = apply_temperature(dev_blend, best_T)
    best_thr = tune_thresholds(cal_cal, y_cal)

    _, _, dev_score = compute_score(y_cal, predict_final(cal_cal, best_thr))
    print(f"\n  Cal score (no leakage): {dev_score:.4f}%")

    full_a = get_loader(X_dev, y_dev, 16, mode='spatial')
    full_b = get_loader(X_dev, y_dev, 16, mode='dct')
    pfa    = get_probs(model_a, full_a)
    pfb    = get_probs(model_b, full_b)
    full_blend = w_a * pfa + w_b * pfb
    full_cal   = apply_temperature(full_blend, best_T)
    full_preds = predict_final(full_cal, best_thr)
    l, b, s    = compute_score(y_dev, full_preds)
    print(f"\n== HONEST FULL DEV SCORE: {s:.4f}%  (lbl={l:.2f}%, bckg={b:.2f}%)")

    retrain_kwargs = dict(
        lr=1.98e-4, weight_decay=0.0203,
        layer_decay=0.859, drop_path=0.190,
        batch_size=16
    )
    model_a_final = retrain_vit(
        X_all, y_all, mode='spatial',
        best_epochs=max(epochs_a, 5),
        model_name="ViT-A", **retrain_kwargs
    )
    model_b_final = retrain_vit(
        X_all, y_all, mode='dct',
        best_epochs=max(epochs_b, 5),
        model_name="ViT-B", **retrain_kwargs
    )

    # Runs both final models on the given features and returns blended, temperature-calibrated probabilities.
    def get_blended_final(X):
        la = get_loader(X, batch_size=16, mode='spatial')
        lb = get_loader(X, batch_size=16, mode='dct')
        pa = get_probs(model_a_final, la)
        pb = get_probs(model_b_final, lb)
        return apply_temperature(w_a * pa + w_b * pb, best_T)

    eval_cal_f = get_blended_final(X_eval)
    eval_preds = predict_final(eval_cal_f, best_thr)

    tr_cal_f = get_blended_final(X_tr)
    tr_preds = predict_final(tr_cal_f, best_thr)
    l, b, s  = compute_score(y_tr, tr_preds)
    print(f"\n== FINAL SCORES ==")
    print(f"  Train: {s:.4f}%  (lbl={l:.2f}%, bckg={b:.2f}%)")
    print(f"  Honest full dev (phase-1 models): reported above")

    print("\nWriting submission files to NoLeak/...")
    write_submission(tr_preds,   tr_idx,   tr_len,   f"{OUT_DIR}/hyp_train_neural.csv")
    write_submission(full_preds, dev_idx,  dev_len,  f"{OUT_DIR}/hyp_dev_neural.csv")
    write_submission(eval_preds, eval_idx, eval_len, f"{OUT_DIR}/hyp_eval_neural.csv")
    print("Done!")


if __name__ == "__main__":
    main()