#!/usr/bin/env python
# train_nn.py
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from train import load_csv, filter_labels, build_features

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# Only use the compact engineered features, not raw DCT or spatial pixels
# build_features returns: [raw_dct (3072) | spatial_flat (3072) | eng features (118)]
# We only want the eng features — last 118 columns
ENG_FEATURES_START = 3072 + 3072   # skip raw dct + spatial


def normalize(Xf):
    mean = Xf.mean(axis=0)
    std  = Xf.std(axis=0) + 1e-8
    return (Xf - mean) / std, mean, std


class PatchDataset(Dataset):
    def __init__(self, Xf_norm, y, augment=True):
        self.Xf      = Xf_norm.astype(np.float32)
        self.y       = y
        self.augment = augment

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

    def __getitem__(self, i):
        x = self.Xf[i].copy()
        if self.augment:
            x += np.random.normal(0, 0.03, x.shape).astype(np.float32)
        return (
            torch.tensor(x,         dtype=torch.float32),
            torch.tensor(self.y[i], dtype=torch.long),
        )


class SimpleNet(nn.Module):
    """3-layer MLP on compact features only."""
    def __init__(self, in_dim, num_classes=9):
        super().__init__()
        self.net = nn.Sequential(
            # Layer 1
            nn.Linear(in_dim, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.3),

            # Layer 2
            nn.Linear(128, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.2),

            # Layer 3 — output
            nn.Linear(64, num_classes),
        )

    def forward(self, x):
        return self.net(x)


def main():
    train_csv, dev_csv, out_file = sys.argv[1], sys.argv[2], sys.argv[3]

    print("Loading data...")
    y_tr, X_tr = load_csv(train_csv)
    y_tr, X_tr = filter_labels(y_tr, X_tr)

    y_dv, X_dv = load_csv(dev_csv)
    y_dv, X_dv = filter_labels(y_dv, X_dv)

    y = np.concatenate([y_tr, y_dv])
    X = np.concatenate([X_tr, X_dv])

    print(f"Samples: {len(y)} (train + dev)")

    print("Building features...")
    Xf_full = build_features(X)

    # Only use compact engineered features
    Xf = Xf_full[:, ENG_FEATURES_START:]
    print(f"Compact feature shape: {Xf.shape}  (dropped raw DCT + spatial pixels)")

    Xf_norm, mean, std = normalize(Xf)

    # Class weights — balanced but moderate
    counts = np.bincount(y, minlength=9).astype(np.float32)
    w = np.zeros(9, dtype=np.float32)
    for cls in np.unique(y):
        w[cls] = 1.0 / counts[cls]
    # Moderate boosts — not as extreme as GB since MLP is more sensitive
    w[5] *= 2.0    # dcis
    w[6] *= 1.8    # indc
    w[3] *= 2.0    # infl
    w[0] *= 1.1    # norm
    w[8] *= 0.8    # bckg
    w = w / w[w > 0].mean()
    print(f"Class weights: { {i: round(float(w[i]),3) for i in np.unique(y)} }")
    class_weights = torch.tensor(w, dtype=torch.float32).to(device)

    dataset = PatchDataset(Xf_norm, y, augment=True)
    loader  = DataLoader(dataset, batch_size=128, shuffle=True, num_workers=0)

    model     = SimpleNet(Xf_norm.shape[1]).to(device)
    criterion = nn.CrossEntropyLoss(weight=class_weights)
    optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

    best_loss  = float('inf')
    best_state = None

    print("Training...")
    for epoch in range(100):
        model.train()
        total_loss = 0
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            loss = criterion(model(xb), yb)
            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            total_loss += loss.item()
        scheduler.step()

        if total_loss < best_loss:
            best_loss  = total_loss
            best_state = {k: v.clone() for k, v in model.state_dict().items()}

        if epoch % 10 == 0:
            print(f"Epoch {epoch:03d}: loss={total_loss:.4f}  lr={scheduler.get_last_lr()[0]:.6f}")

    print(f"Best loss: {best_loss:.4f}")
    model.load_state_dict(best_state)

    torch.save({
        'state_dict':        model.state_dict(),
        'mean':              mean,
        'std':               std,
        'in_dim':            Xf_norm.shape[1],
        'eng_features_start': ENG_FEATURES_START,
    }, out_file)
    print("Saved.")


if __name__ == "__main__":
    main()