# %% Импорт библиотек

import time
import torch
import kagglehub
import pandas as pd
import torch.nn as nn
import pytorch_lightning as pl

from PIL import Image
from torch import optim
from pathlib import Path
from torchmetrics import MetricCollection
from typing import Optional, Callable, Any
from torchvision import datasets, transforms
from torchmetrics.classification import Accuracy
from torch.utils.data import DataLoader, Dataset

# Логирование данных
from pytorch_lightning.loggers import MLFlowLogger
# автоматическое отслеживание lr
from pytorch_lightning.callbacks import LearningRateMonitor
# ранние остановки
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
# сохранение весов
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
# %% Параметры

dataset_root = "/media/daniel/developments/Datasets/ImageNet"
batch_size = 128
num_workers = 8
num_epoch = 90
learning_rate = 3e-4

# %% Трансформации

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),

    transforms.ToTensor(),

    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

eval_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),

    transforms.ToTensor(),

    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# %% Скачивание ImageNet

path = kagglehub.competition_download('imagenet-object-localization-challenge',
                                      output_dir=dataset_root,
                                      force_download=False)
print("Path to competition files:", path)
# %% Датасет

TransformType = Optional[Callable[..., Any]]

class ImageNetValDataset(Dataset):
    def __init__(
        self,
        image_dir,
        labels_file,
        class_to_idx,
        classes,
        transform=None
    ):
        self.image_dir = Path(image_dir)
        self.transform = transform

        self.class_to_idx = class_to_idx
        self.classes = classes

        df = pd.read_csv(labels_file)

        self.samples = []

        for _, row in df.iterrows():
            image_id = row["ImageId"]
            synset = row["PredictionString"].split()[0]
            label = self.class_to_idx[synset]
            image_path = self.image_dir / f"{image_id}.JPEG"
            self.samples.append(
                (image_path, label)
            )

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

    def __getitem__(self, idx):
        image_path, label = self.samples[idx]

        image = Image.open(image_path).convert("RGB")

        if self.transform is not None:
            image = self.transform(image)

        return image, label


class ImageNetDataModule(pl.LightningDataModule):
    def __init__(self, 
                 dataset_root: Path,
                 train_transform: TransformType,
                 eval_transform: TransformType,
                 batch_size: int = 64,
                 num_workers:int = 8
                 ):
        super().__init__()
        
        self.dataset_root = Path(dataset_root)
        
        self.train_transform = train_transform
        self.eval_transform = eval_transform
        self.num_workers = num_workers
        self.batch_size = batch_size
        
        self.imagenet_root = Path(dataset_root) / "ILSVRC"
        
        self.train_dir = self.imagenet_root / "Data" / "CLS-LOC" / "train"
        self.val_dir = self.imagenet_root / "Data" / "CLS-LOC" / "val"
        
        self.mapping_file = Path(dataset_root) / "LOC_synset_mapping.txt"
        self.val_labels_file = Path(dataset_root) / "LOC_val_solution.csv"

        self.train_dataset = None
        self.val_dataset = None
        
        self.class_to_idx = None
        self.classes = None
        
    def setup(self, stage=None):
        synset_to_name = {}
        with open(self.mapping_file, "r", encoding="utf-8") as f:
            for line in f:
                synset, class_name = line.strip().split(" ", 1)
                synset_to_name[synset] = class_name
        
        train_dataset = datasets.ImageFolder(root=self.train_dir, 
                                             transform=self.train_transform)
    
        self.class_to_idx = train_dataset.class_to_idx.copy()
        self.classes = [synset_to_name[synset] 
                        for synset in train_dataset.classes]
        train_dataset.classes = self.classes
        self.train_dataset = train_dataset
        
        self.val_dataset = ImageNetValDataset(
            image_dir=self.val_dir,
            labels_file=self.val_labels_file,
            class_to_idx=self.class_to_idx,
            classes=self.classes,
            transform=self.eval_transform
        )
        
    def train_dataloader(self):
        return DataLoader(
            dataset=self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=self.num_workers > 0,
            prefetch_factor=4 if self.num_workers > 0 else None
        )
    
    def val_dataloader(self):
        return DataLoader(
            dataset=self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=self.num_workers > 0,
            prefetch_factor=4 if self.num_workers > 0 else None
        )
    
    def test_dataloader(self):
        return DataLoader(
            dataset=self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=self.num_workers > 0    
        )

# %% AlexNet

class AlexNetCustom(nn.Module):
    def __init__(self,
                 in_channels: int = 3,
                 out_classes: int=1_000):
        super().__init__()
        
        self._out_classes = out_classes
        
        self.feature = nn.Sequential(
            nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            
            nn.Conv2d(96, 256, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            
            nn.Conv2d(256, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(384, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2)
        )
        
        # Добавлено в современной реализации
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) 
        
        self.mlp = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(256*6*6, 4096),
            nn.ReLU(inplace=True),
            
            nn.Dropout(p=0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            
            nn.Linear(4096, self._out_classes)
        )
        
    def forward(self, image):
        x = self.feature(image)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.mlp(x)
        
        return x
    
# %% Классификатор обучатор

class ClassificationLightning(pl.LightningModule):
    def __init__(self,
                 model: nn.Module,
                 num_classes: int,
                 optimizer_class,
                 optimizer_kwargs: dict,
                 scheduler_class=None,
                 scheduler_kwargs: dict | None = None,
                 scheduler_monitor: str = "val/loss"):
        super().__init__()
        
        
        scheduler_kwargs = scheduler_kwargs or {}
        
        self.model = model
        self.optimizer_class = optimizer_class
        self.scheduler_class = scheduler_class
        
        num_parameters = sum(
            p.numel()
            for p in model.parameters()
        )
        
        num_trainable_parameters = sum(
            p.numel()
            for p in model.parameters()
            if p.requires_grad
        )
        
        self.save_hyperparameters({
            "model": model.__class__.__name__,
            "num_classes": num_classes,
        
            "num_parameters": num_parameters,
            "num_trainable_parameters": num_trainable_parameters,
        
            "optimizer": optimizer_class.__name__,
            "optimizer_kwargs": optimizer_kwargs,
        
            "scheduler": (
                scheduler_class.__name__
                if scheduler_class is not None
                else None
            ),
            "scheduler_kwargs": scheduler_kwargs,
            "scheduler_monitor": scheduler_monitor
        })
        
        
        self.criterion = nn.CrossEntropyLoss()
        self.train_metrics = MetricCollection({
            "acc_top1": Accuracy(task="multiclass", 
                                 num_classes=self.hparams.num_classes,
                                 top_k=1),
            "acc_top5": Accuracy(task="multiclass", 
                                 num_classes=self.hparams.num_classes,
                                 top_k=5),  
        })
        
        self.val_metrics = self.train_metrics.clone()
        self.test_metrics = self.train_metrics.clone()
        
    def forward(self, image):
        return self.model(image)
    
    def _shared_step(self,
                    sample,
                    stage,
                    metrics):
        x, y = sample
        
        logits = self(x)
        loss = self.criterion(logits, y)
        
        metrics.update(logits, y)
        
        self.log(f"{stage}/loss", 
                 loss, 
                 prog_bar=True, 
                 on_epoch=True, 
                 on_step=False
         )
        
        return loss
    
    def _log_metrics(self, stage, metrics):
        computed = metrics.compute()
        
        for name, value in computed.items():
            self.log(f"{stage}/{name}", value, prog_bar=True)
        
        metrics.reset()
    
    def on_train_epoch_start(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        
        self._epoch_start_time = time.perf_counter()
    
    def training_step(self, batch, batch_idx):
        return self._shared_step(batch, "train", self.train_metrics)
    
    def on_train_epoch_end(self):
        self._log_metrics("train", self.train_metrics)
        
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        
        epoch_time = time.perf_counter() - self._epoch_start_time
        lr = self.optimizers().param_groups[0]["lr"]
        
        self.log("train/lr", lr, prog_bar=True)
        self.log("train/epoch_time", epoch_time)
    
    def validation_step(self, batch, batch_idx):
        return self._shared_step(batch, "val", self.val_metrics)
    
    def on_validation_epoch_end(self):
        self._log_metrics("val", self.val_metrics)
    
    def test_step(self, batch, batch_idx):
        return self._shared_step(batch, "test", self.test_metrics)
    
    def on_test_epoch_end(self):
       self._log_metrics("test", self.test_metrics)
    
    def configure_optimizers(self):
        optimizer = self.optimizer_class(
            self.parameters(),
            **self.hparams.optimizer_kwargs
        )
        
        if self.scheduler_class is None:
            return optimizer
        
        scheduler = self.scheduler_class(
            optimizer,
            **self.hparams.scheduler_kwargs
        )
        
        scheduler_config = {
            "scheduler": scheduler,
            "interval": "epoch",
            "frequency": 1
        }
        
        if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
            scheduler_config["monitor"] = self.hparams.scheduler_monitor
        
        return {
            "optimizer": optimizer,
            "lr_scheduler": scheduler_config
        }

# %% Обучение

pl.seed_everything(42, workers=True)
torch.set_float32_matmul_precision("high")

imagenet_dm = ImageNetDataModule(
    dataset_root=dataset_root,
    train_transform=train_transform,
    eval_transform=eval_transform,
    batch_size=batch_size,
    num_workers=num_workers    
)


optimizer_Adam = optim.Adam
optimizer_Adam_kwargs = {
    "lr": 3e-4,
    "weight_decay": 1e-4
}

scheduler_ReduceLR = optim.lr_scheduler.ReduceLROnPlateau
scheduler_ReduceLR_kwargs={
    "mode": "min",
    "factor": 0.1,
    "patience": 5
}

callbacks = [
    ModelCheckpoint(
        dirpath="models",
        filename="epoch={epoch}",
        auto_insert_metric_name=False,
        save_last=True,
        save_top_k=1,
        monitor="val/acc_top1",
        mode="max"
    ),

    LearningRateMonitor(
        logging_interval="epoch"
    ),

    EarlyStopping(
        monitor="val/loss",
        min_delta=2e-4,
        patience=10,
        verbose=False,
        mode="min"
    )
]

logger = MLFlowLogger(
    experiment_name="ImageNetMyAlgoritms",
    run_name="AlexNetCustom",
    tracking_uri="http://10.16.18.1:8055",
    log_model=True,
) 

alexnet = AlexNetCustom(in_channels=3, out_classes=1000)
model_adam = ClassificationLightning(
    model=alexnet,
    num_classes=1000,
    optimizer_class=optimizer_Adam,
    optimizer_kwargs=optimizer_Adam_kwargs,
    scheduler_class=scheduler_ReduceLR,
    scheduler_kwargs=scheduler_ReduceLR_kwargs,
    scheduler_monitor="val/loss"    
)

trainer = pl.Trainer(
    max_epochs=num_epoch,
    accelerator="gpu",
    devices=1,
    logger=logger,
    # precision="16-mixed",
    callbacks=callbacks
)

trainer.fit(model=model_adam,
            datamodule=imagenet_dm)