import lightning as L
import torch
import torch.nn as nn
import torch.nn.functional as F
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.strategies import FSDP2Strategy
from torch.utils.data import DataLoader, TensorDataset


class SimpleModel(L.LightningModule):
    def __init__(self, input_dim=32, hidden_dim=64, output_dim=10):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.layer2 = nn.Linear(hidden_dim, hidden_dim)
        self.layer3 = nn.Linear(hidden_dim, output_dim)
        self.save_hyperparameters()
        
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        return self.layer3(x)
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        self.log("train_loss", loss, prog_bar=True)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        self.log("val_loss", loss, prog_bar=True)
        
    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        self.log("test_loss", loss, prog_bar=True)
        
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)


class RandomDataModule(L.LightningDataModule):
    def __init__(self, input_dim=32, output_dim=10, num_samples=1000, batch_size=32):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.num_samples = num_samples
        self.batch_size = batch_size
        
    def setup(self, stage=None):
        # Generate random data
        if stage == 'fit' or stage is None:
            # Training data
            x_train = torch.randn(self.num_samples, self.input_dim)
            y_train = torch.randint(0, self.output_dim, (self.num_samples,))
            self.train_dataset = TensorDataset(x_train, y_train)
            
            # Validation data
            x_val = torch.randn(self.num_samples // 5, self.input_dim)
            y_val = torch.randint(0, self.output_dim, (self.num_samples // 5,))
            self.val_dataset = TensorDataset(x_val, y_val)
            
        if stage == 'test' or stage is None:
            # Test data
            x_test = torch.randn(self.num_samples // 5, self.input_dim)
            y_test = torch.randint(0, self.output_dim, (self.num_samples // 5,))
            self.test_dataset = TensorDataset(x_test, y_test)
    
    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)
    
    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size)
    
    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size)


def main():
    # Set up model and data
    input_dim = 32
    hidden_dim = 64
    output_dim = 10
    
    with torch.device("meta"):
        model = SimpleModel(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim)
    data_module = RandomDataModule(input_dim=input_dim, output_dim=output_dim)
    
    # Set up the checkpoint callback to save the best model
    checkpoint_callback = ModelCheckpoint(
        monitor="val_loss",
        dirpath="checkpoints/",
        filename="model-{epoch:02d}-{val_loss:.2f}",
        save_top_k=1,
        mode="min",
    )

    # # Set up the FSDP strategy
    # fsdp_strategy = FSDP2Strategy(
    # )
    
    # Set up the trainer
    trainer = L.Trainer(
        max_epochs=10,
        callbacks=[checkpoint_callback],
        strategy="fsdp2",
        accelerator="auto",  # Automatically detect available accelerator
        devices="auto",      # Use all available devices
    )
    
    # Train the model
    trainer.fit(model, data_module)
    
    if trainer.is_global_zero:
        print(f"Done training! Best model saved at: {checkpoint_callback.best_model_path}")


if __name__ == "__main__":
    main()