import torch
import torch.nn as nn
import pandas as pd
import numpy as np
from sklearn.metrics import roc_auc_score, accuracy_score, f1_score, confusion_matrix
import pysam
from tqdm import tqdm

device = "cpu"

MODEL_PATH = "data/mutation_predictor_splice_v4.pt"
BENCH_PATH = "external_validation/clinvar_splice_micro_benchmark_seq.csv"
FASTA_PATH = "external_validation/reference/Homo_sapiens.GRCh38.dna.primary_assembly.fa"

WIN = 401
HALF = 200
SEQ_LEN = 399

print("Loading FASTA...")
fasta = pysam.FastaFile(FASTA_PATH)

print("Loading benchmark...")
df = pd.read_csv(BENCH_PATH)

# ==========================================================
# Encoding utilities
# ==========================================================

BASE_MAP = {"A":0,"C":1,"G":2,"T":3}
COMP = {"A":"T","T":"A","C":"G","G":"C","N":"N"}

MUT_TYPES = [
    "A>C","A>G","A>T",
    "C>A","C>G","C>T",
    "G>A","G>C","G>T",
    "T>A","T>C","T>G"
]

def encode_seq(seq):
    arr = np.zeros((11, SEQ_LEN), dtype=np.float32)
    half = WIN // 2

    for i in range(SEQ_LEN):
        j = i + 1
        b = seq[j] if j < len(seq) else "N"

        if b in BASE_MAP:
            arr[BASE_MAP[b], i] = 1
            cb = COMP[b]
            if cb in BASE_MAP:
                arr[4 + BASE_MAP[cb], i] = 1

        arr[8, i] = (j - half) / half
        arr[9, i] = 1 if seq[j:j+2] == "GT" else 0
        arr[10, i] = 1 if seq[j:j+2] == "AG" else 0

    return arr


def mut_onehot(ref, alt):
    key = f"{ref.upper()}>{alt.upper()}"
    v = np.zeros(12, dtype=np.float32)
    if key in MUT_TYPES:
        v[MUT_TYPES.index(key)] = 1.0
    return v


def region_flags(splice_dist):
    try:
        d = float(splice_dist)
    except:
        d = 999.0

    return np.array([
        1.0 if d > 10 else 0.0,
        1.0 if 2 < d <= 10 else 0.0,
    ], dtype=np.float32)


def splice_flags(splice_dist, splice_motif):
    try:
        d = float(splice_dist)
    except:
        d = 999.0

    m = str(splice_motif).upper()

    return np.array([
        1.0 if m in ["GT","GU"] else 0.0,
        1.0 if m == "AG" else 0.0,
        1.0 if d <= 8 else 0.0,
    ], dtype=np.float32)


# ==========================================================
# Model definition (MATCHES TRAINING EXACTLY)
# ==========================================================

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv1d(11,64,7,padding=3)
        self.conv2 = nn.Conv1d(64,128,5,padding=2)
        self.conv3 = nn.Conv1d(128,256,3,padding=1)

        self.pool = nn.AdaptiveAvgPool1d(1)

        self.mut_fc = nn.Linear(12,32)
        self.region_fc = nn.Linear(2,8)
        self.splice_fc = nn.Linear(3,16)

        self.fc1 = nn.Linear(312,128)
        self.fc2 = nn.Linear(128,64)
        self.fc3 = nn.Linear(64,1)

        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)

    def forward(self,seq,mut,region,splice):

        s = self.relu(self.conv1(seq))
        s = self.relu(self.conv2(s))
        s = self.relu(self.conv3(s))
        s = self.pool(s).squeeze(-1)

        m = self.relu(self.mut_fc(mut))
        r = self.relu(self.region_fc(region))
        sp = self.relu(self.splice_fc(splice))

        x = torch.cat([s,m,r,sp],dim=1)

        x = self.dropout(self.relu(self.fc1(x)))
        x = self.relu(self.fc2(x))

        return self.fc3(x)

model = Model().to(device)

ckpt = torch.load(MODEL_PATH, map_location=device)
model.load_state_dict(ckpt)
model.eval()

# ==========================================================
# Inference
# ==========================================================

probs = []
labels = []

print("Running inference...")

with torch.no_grad():
    for _, row in tqdm(df.iterrows(), total=len(df)):

        chrom = str(row["chrom"])
        pos = int(row["pos"])

        start = pos - 1 - HALF
        end = pos - 1 + HALF + 1

        seq = fasta.fetch(chrom, start, end).upper()

        seq_enc = encode_seq(seq)
        mut_enc = mut_onehot(row["ref"], row["alt"])
        reg_enc = region_flags(row.get("splice_dist", 999))
        spl_enc = splice_flags(row.get("splice_dist", 999),
                               row.get("splice_motif","."))

        seq_t = torch.tensor(seq_enc).unsqueeze(0)
        mut_t = torch.tensor(mut_enc).unsqueeze(0)
        reg_t = torch.tensor(reg_enc).unsqueeze(0)
        spl_t = torch.tensor(spl_enc).unsqueeze(0)

        logit = model(seq_t, mut_t, reg_t, spl_t)
        p = torch.sigmoid(logit).item()

        probs.append(p)
        labels.append(row["true_label"])

# ==========================================================
# Metrics
# ==========================================================

auc = roc_auc_score(labels, probs)

pred = [1 if p >= 0.5 else 0 for p in probs]

acc = accuracy_score(labels, pred)
f1 = f1_score(labels, pred)
tn, fp, fn, tp = confusion_matrix(labels, pred).ravel()

sens = tp / (tp + fn)
spec = tn / (tn + fp)

print("\n===== EXTERNAL RESULTS (v4) =====")
print(f"AUC : {auc:.4f}")
print(f"ACC : {acc:.4f}")
print(f"F1  : {f1:.4f}")
print(f"SENS: {sens:.4f}")
print(f"SPEC: {spec:.4f}")
print(f"TP={tp}  TN={tn}  FP={fp}  FN={fn}")


# ==========================================================
# Save predictions for comparison
# ==========================================================

import pandas as pd

out = pd.DataFrame({
    "prob_v4": probs,
    "true_label": labels
})

out.to_csv("external_validation/v4_predictions.csv", index=False)
print("Saved external_validation/v4_predictions.csv")