import os
import torch
import mlflow
import pytorch_lightning as pl

from torch import optim
from dotenv import load_dotenv
from torchvision.transforms import v2
from datasets import ImageNetDataModule
from train import ClassificationLightning
from custom_arhitecture import VGGCustom

from pytorch_lightning.loggers import MLFlowLogger
from pytorch_lightning.callbacks import LearningRateMonitor
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint


load_dotenv()

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

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

train_transform = v2.Compose([
    v2.RandomResizedCrop(
        size=(224, 224),
        scale=(0.08, 1.0),
        ratio=(3 / 4, 4 / 3),
        antialias=True
    ),
    v2.RandomHorizontalFlip(p=0.5),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(
        mean=IMAGENET_MEAN,
        std=IMAGENET_STD
    )
])

eval_transform = v2.Compose([
    v2.Resize(256, antialias=True),
    v2.CenterCrop(224),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(
        mean=IMAGENET_MEAN,
        std=IMAGENET_STD
    )
])


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_kwargs = {
    "lr": 3e-4,
    "weight_decay": 1e-4
}

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"
    )
]

# Системные метрики
mlflow.enable_system_metrics_logging()
mlflow.set_system_metrics_sampling_interval(5)

logger = MLFlowLogger(
    experiment_name=os.getenv("EXPERIMENT_NAME"),
    run_name="VGG16Custom",
    tracking_uri=os.getenv("TRACKING_URI"),
    log_model=True,
) 


vgg = VGGCustom(architecture="VGG16",
                num_channels=3,
                num_classes=1000,
                batch_normalization=False,
                dropout=.5)

# Необходимо чтобы отрпаботал LazyLinear и принял размерности!
# with torch.no_grad():
#     alexnet(torch.zeros(1, 3, 224, 224))

model = ClassificationLightning(model=vgg,
                                num_classes=1000,
                                optimizer_class=optim.Adam,
                                optimizer_kwargs=optimizer_Adam_kwargs,
                                scheduler_class=optim.lr_scheduler.ReduceLROnPlateau,
                                scheduler_kwargs=scheduler_ReduceLR_kwargs,
                                scheduler_monitor="val/loss")

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

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