import torch
from models import MLP, Resnet, VGG
# import torchvision.models as models
import torch.nn as nn
import torch.optim as optim
import os
import csv
import shutil
import get_dataset
import argparse
from twain import *
import lam


def collect_ternary_stats(model):
    """Collect ternary weight statistics from each TnLinear layer."""
    stats = []
    for name, module in model.named_modules():
        if isinstance(module, MLP.TnLinear):
            with torch.no_grad():
                w_tn = MLP.linear_Ternarize(module.weight.data)
                n = w_tn.numel()
                p = (w_tn == -1).float().sum().item() / n
                q = (w_tn == 0).float().sum().item() / n
                pos = 1.0 - p - q
            stats.append({
                'layer_name': name,
                'p(-1)': p,
                'q(0)': q,
                'pos(+1)': pos,
                'numel': n,
            })
    return stats


def append_stats_csv(csv_path, round_idx, acc, stats):
    file_exists = os.path.exists(csv_path)
    with open(csv_path, 'a', newline='') as f:
        writer = csv.writer(f)
        if not file_exists:
            writer.writerow(['round', 'accuracy', 'layer_name', 'p(-1)', 'q(0)', 'pos(+1)', 'numel'])
        for s in stats:
            writer.writerow([round_idx, acc, s['layer_name'], s['p(-1)'], s['q(0)'], s['pos(+1)'], s['numel']])


def Run(folder, model, train_set, test_set, lr=0.2, lr_decay=0.5, epochs=10, save_point=1, rounds=2, pretrained_path=None):
    device = 'cuda:0'
    model.to(device)

    # Load pretrained weights if provided
    if pretrained_path and os.path.exists(pretrained_path):
        model.load_state_dict(torch.load(pretrained_path))
        print(f"Loaded pretrained weights from {pretrained_path}")

    criterion = nn.CrossEntropyLoss()

    if not os.path.exists(folder):
        os.makedirs(folder)
    csv_path = os.path.join(folder, 'ternary_dist.csv')
    for i in range(rounds):
        optimizer = optim.SGD(model.parameters(), lr=lr * (lr_decay ** i), momentum=0.9, weight_decay=5e-4)
        for j in range(save_point):
            acc = train(model, train_set, test_set, device, criterion, optimizer, epochs)
        torch.save(model.state_dict(), f'{folder}/{i}_{acc}.pth')
        stats = collect_ternary_stats(model)
        append_stats_csv(csv_path, i, acc, stats)


if __name__ == '__main__':
    # Resolve relative paths from the script directory.
    os.chdir(os.path.dirname(os.path.abspath(__file__)))

    parser = argparse.ArgumentParser()
    parser.add_argument('--use-moma', dest='use_moma', action='store_true',
                        help='Enable MOMA neighborhood thresholding (default).')
    parser.add_argument('--no-moma', dest='use_moma', action='store_false',
                        help='Disable MOMA and use the mean-based threshold.')
    parser.set_defaults(use_moma=True)
    args = parser.parse_args()
    lam.use_moma = args.use_moma

    train_loader = get_dataset.get_mnist_dataloader('mnist_data', batch_size=512, num_workers=0, spilt='train')
    test_loader = get_dataset.get_mnist_dataloader('mnist_data', batch_size=512, num_workers=0, spilt='test')

    # Define lambda values
    lam_values = [
                #   0.001,
                #   0.01,
                  0.1,
                #   1,
                #   10
                ]
    c_list=[
            # 0.1,
            # 0.3,
            # 0.5,
            # 0.7,
            0.9
            ]
    # Path to the pretrained model
    pretrained_model_path = "fp.pth"
    method_tag = 'moma' if lam.use_moma else 'mean'
    for c in c_list:
        for x in lam_values:
            lam.lamda = x
            lam.c_val=c
            if lam.use_moma:
                name = f"mlp_tn_moma(c={c})/lam{x}"
            else:
                name = f"mlp_tn/run"
            tn_model = [[name, MLP.MLP2(784, 512, 10)]]
            for name, m in tn_model:
                print(m)
                MLP.l()
                print(f"[method] {method_tag}")
                Run(name, m, train_loader, test_loader, pretrained_path=pretrained_model_path)
