# ============================================================
# MutationPredictorCNN_v4 Training Script (401 bp FASTA)
# Proper sequence-based training
# ============================================================

import argparse
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import roc_auc_score
import pysam
from tqdm import tqdm
import os

# ============================================================
# Arguments
# ============================================================

parser = argparse.ArgumentParser()

parser.add_argument("--train_csv", required=True)
parser.add_argument("--fasta", required=True)
parser.add_argument("--output_model", required=True)

parser.add_argument("--epochs", type=int, default=30)
parser.add_argument("--batch_size", type=int, default=256)
parser.add_argument("--num_workers", type=int, default=8)
parser.add_argument("--lr", type=float, default=0.001)

args = parser.parse_args()

# ============================================================
# Config
# ============================================================

WINDOW = 401
HALF = WINDOW // 2
SEQ_LEN = WINDOW - 2
DEVICE = "cpu"

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

# ============================================================
# Encoding
# ============================================================

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

def fetch_seq(chrom, pos):

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

    try:
        return fasta.fetch(str(chrom), start, end).upper()
    except:
        try:
            return fasta.fetch("chr"+str(chrom), start, end).upper()
        except:
            return None


def encode_seq(seq):

    arr = np.zeros((11, SEQ_LEN), dtype=np.float32)

    for i in range(SEQ_LEN):

        j = i + 1
        base = seq[j] if j < len(seq) else "N"

        if base in BASE_MAP:
            arr[BASE_MAP[base], i] = 1
            comp = COMP[base]
            if comp in BASE_MAP:
                arr[4 + BASE_MAP[comp], i] = 1

        arr[8, i] = (j - HALF) / HALF

        if seq[j:j+2] == "GT":
            arr[9, i] = 1

        if seq[j:j+2] == "AG":
            arr[10, i] = 1

    return arr


def mut_onehot(ref, alt):

    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"
    ]

    vec = np.zeros(12, dtype=np.float32)

    key = f"{ref}>{alt}"

    if key in types:
        vec[types.index(key)] = 1

    return vec


# ============================================================
# Dataset
# ============================================================

class SpliceDataset(Dataset):

    def __init__(self, df):

        self.df = df.reset_index(drop=True)

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):

        row = self.df.iloc[idx]

        seq = fetch_seq(row.chrom, int(row.pos))

        if seq is None or len(seq) != WINDOW:
            seq = "N" * WINDOW

        seq_enc = encode_seq(seq)

        mut = mut_onehot(row.ref, row.alt)

        region = np.zeros(2, dtype=np.float32)
        splice = np.zeros(3, dtype=np.float32)

        label = float(row.label)

        return (
            torch.tensor(seq_enc),
            torch.tensor(mut),
            torch.tensor(region),
            torch.tensor(splice),
            torch.tensor(label)
        )


# ============================================================
# Model
# ============================================================

class MutationPredictorCNN_v4(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):

        x = self.relu(self.conv1(seq))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))

        x = self.pool(x).squeeze(-1)

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

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

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

        return self.fc3(x)


# ============================================================
# Load dataset
# ============================================================

print("Loading dataset...")

df = pd.read_csv(args.train_csv)

train_ds = SpliceDataset(df)

train_dl = DataLoader(
    train_ds,
    batch_size=args.batch_size,
    shuffle=True,
    num_workers=args.num_workers
)

# ============================================================
# Train
# ============================================================

model = MutationPredictorCNN_v4().to(DEVICE)

criterion = nn.BCEWithLogitsLoss()

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=args.lr
)

best_auc = 0

for epoch in range(args.epochs):

    model.train()

    losses = []
    probs = []
    labels = []

    for seq, mut, region, splice, label in train_dl:

        seq = seq.to(DEVICE)
        mut = mut.to(DEVICE)
        region = region.to(DEVICE)
        splice = splice.to(DEVICE)
        label = label.to(DEVICE).unsqueeze(1)

        optimizer.zero_grad()

        logits = model(seq, mut, region, splice)

        loss = criterion(logits, label)

        loss.backward()

        optimizer.step()

        losses.append(loss.item())

        probs.extend(torch.sigmoid(logits).detach().cpu().numpy())
        labels.extend(label.cpu().numpy())

    auc = roc_auc_score(labels, probs)

    print(f"Epoch {epoch+1}/{args.epochs} Loss={np.mean(losses):.4f} AUC={auc:.4f}")

    if auc > best_auc:

        best_auc = auc

        torch.save(
            model.state_dict(),
            args.output_model
        )

        print("Saved best model")

print("Training complete.")