import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import json
import os

from model import PhysicsGuidedBNN
from dataset import WindTurbineDataset

def train_model(config_path: str, data_path: str, save_dir: str = "./checkpoint"):
    with open(config_path, "r") as f:
        config = json.load(f)

    os.makedirs(save_dir, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    dataset = WindTurbineDataset(data_path, is_train=True)
    dataloader = DataLoader(dataset, batch_size=config["batch_size"], shuffle=True)

    model = PhysicsGuidedBNN(config).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=config["learning_rate"], weight_decay=config["weight_decay"])
    ce_criterion = nn.CrossEntropyLoss()

    model.train()
    print(f"Starting training on device: {device}")

    for epoch in range(config["epochs"]):
        total_loss = 0.0
        total_ce = 0.0
        total_kl = 0.0
        total_phys = 0.0

        for batch_x, batch_y in dataloader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            optimizer.zero_grad()

            logits = model(batch_x)
            
            # Loss Components
            loss_ce = ce_criterion(logits, batch_y)
            loss_kl = model.total_kl_divergence() * config["beta_kl"]
            loss_physics = model.compute_physics_loss(batch_x, logits) * config["lambda_physics"]

            # Combined Loss
            loss = loss_ce + loss_kl + loss_physics
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
            total_ce += loss_ce.item()
            total_kl += loss_kl.item()
            total_phys += loss_physics.item()

        if (epoch + 1) % 5 == 0 or epoch == 0:
            print(f"Epoch [{epoch+1}/{config['epochs']}] | Loss: {total_loss/len(dataloader):.4f} | "
                  f"CE: {total_ce/len(dataloader):.4f} | KL: {total_kl/len(dataloader):.4f} | "
                  f"Physics: {total_phys/len(dataloader):.4f}")

    # Save model weights and configuration
    torch.save(model.state_dict(), os.path.join(save_dir, "pytorch_model.bin"))
    with open(os.path.join(save_dir, "config.json"), "w") as f:
        json.dump(config, f, indent=2)

    print(f"Model saved successfully to {save_dir}")

if __name__ == "__main__":
    # Example execution call
    train_model("config.json", "scada_sample.csv")