import numpy as np
from sklearn.preprocessing import StandardScaler
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

BATCH_SIZE = 128
EPOCHS = 500
LEARNING_RATE = 0.001
# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DEVICE = torch.device("cpu")
CLASS_MAP = {0: 0, 2: 1, 3: 2, 5: 3, 6: 4, 8: 5}
INV_MAP = {v: k for k, v in CLASS_MAP.items()}

def load_csv(path):
    data = np.loadtxt(path, delimiter=",")
    y = data[:, 0].astype(int)
    X = data[:, 1:]
    return X, y

def save_csv(filename, labels):
    with open(filename, "w") as f:
        f.write("label\n")
        for lbl in labels:
            f.write(f"{int(lbl)}\n")

def preprocess_for_cnn(X, y, filter_ignored=False):
    if filter_ignored:
        mask = np.isin(y, list(CLASS_MAP.keys()))
        X, y = X[mask], y[mask]
        y = np.array([CLASS_MAP[val] for val in y])
    
    # Reshape (Batch, 3072) -> (Batch, 3, 32, 32)
    X_reshaped = X.reshape(-1, 3, 32, 32)
    return torch.FloatTensor(X_reshaped), torch.LongTensor(y)

def augment_dct_batch(batch_X):
    """
    Applies augmentations directly to the DCT tensors.
    """
    # 1. Random Horizontal/Vertical Flips
    if np.random.rand() > 0.5:
        batch_X = torch.flip(batch_X, [3]) # Horizontal
    if np.random.rand() > 0.5:
        batch_X = torch.flip(batch_X, [2]) # Vertical
        
    # 2. Add Frequency Jitter (Gaussian Noise)
    # This prevents overfitting to exact DCT values
    noise = torch.randn_like(batch_X) * 0.05 
    batch_X = batch_X + noise
    
    return batch_X

class DCT_CNN(nn.Module):
    def __init__(self):
        super(DCT_CNN, self).__init__()
        self.conv_layers = nn.Sequential(
            # Fewer filters (32 -> 16) to prevent memorizing fine noise
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(16),
            nn.MaxPool2d(2), # 16x16
            
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2)  # 8x8
        )
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            # Smaller hidden layer (2048 -> 128)
            nn.Linear(32 * 8 * 8, 128),
            nn.ReLU(),
            nn.Dropout(0.6), # Aggressive dropout to fight 99% correlation
            nn.Linear(128, 6)
        )

    def forward(self, x):
        return self.fc_layers(self.conv_layers(x))

if __name__ == "__main__":
    # Parameters
    test = False

    # 1. Load Data
    print("LOADING DATA...")
    X_train, y_train = load_csv("data/train.csv")
    X_dev, y_dev = load_csv("data/dev.csv")
    X_eval, y_eval = load_csv("data/eval.csv")

    if test:
        print("RUNNING IN TEST MODE: Training on Train set only.")
        X_train_pool = X_train
        y_train_pool = y_train
    else:
        print("RUNNING IN FINAL MODE: Training on Combined (Train + Dev) set.")
        X_train_pool = np.vstack((X_train, X_dev))
        y_train_pool = np.concatenate((y_train, y_dev))

    scaler = StandardScaler()
    X_train_pool = scaler.fit_transform(X_train_pool)
    X_train = scaler.transform(X_train)
    X_dev = scaler.transform(X_dev)
    X_eval = scaler.transform(X_eval)

    X_train_pool, y_train_pool = preprocess_for_cnn(X_train_pool, y_train_pool, filter_ignored=True)
    X_train, y_train = preprocess_for_cnn(X_train, y_train, filter_ignored=False)
    X_dev, y_dev = preprocess_for_cnn(X_dev, y_dev, filter_ignored=False)
    X_eval, y_eval = preprocess_for_cnn(X_eval, y_eval, filter_ignored=False)

    train_loader = DataLoader(TensorDataset(X_train_pool, y_train_pool), batch_size=BATCH_SIZE, shuffle=True)

    print("Training Model")
    model = DCT_CNN().to(DEVICE)
    class_weights = torch.FloatTensor([2.0, 1.0, 1.0, 3.0, 3.0, 3.0]).to(DEVICE)
    criterion = nn.CrossEntropyLoss(weight=class_weights, label_smoothing=0.2)
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
    model.train()
    for epoch in range(EPOCHS):
        total_loss = 0
        for batch_X, batch_y in train_loader:
            batch_X, batch_y = augment_dct_batch(batch_X.to(DEVICE)), batch_y.to(DEVICE)
            # batch_X, batch_y = batch_X.to(DEVICE), batch_y.to(DEVICE)
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        if (epoch+1) % 5 == 0:
            print(f"Epoch {epoch+1}/{EPOCHS}, Loss: {total_loss/len(train_loader):.4f}")

    print("Evaluate Model")
    model.eval()
    def predict(tensor_x):
        with torch.no_grad():
            outputs = model(tensor_x.to(DEVICE))
            _, predicted_idx = torch.max(outputs, 1)
            # Map indices back to original labels (0, 2, 3, 5, 6, 8)
            return [INV_MAP[idx.item()] for idx in predicted_idx]
        
    y_train_pred = predict(X_train)   
    y_dev_pred   = predict(X_dev) 
    y_eval_pred  = predict(X_eval)  

    save_csv("file/ref_train.csv", y_train)
    save_csv("file/ref_dev.csv", y_dev)
    save_csv("file/ref_eval.csv", y_eval)
    save_csv("file/hyp_train_nn.csv", y_train_pred)
    save_csv("file/hyp_dev_nn.csv", y_dev_pred)
    save_csv("file/hyp_eval_nn.csv", y_eval_pred)