from __future__ import annotations
import json
from pathlib import Path

import pytorch_lightning as pl
import timm
import torch
import torch.nn as nn
import torchmetrics
from pydantic import BaseModel, Field

from realfake.data import get_dss, get_dls
from realfake.utils import Args

N_CLASSES = 2


class AcceleratorParams(BaseModel):
    """PyTorch Lightning accelerator parameters."""

    name: str = Field("gpu")
    devices: int = Field(4)
    strategy: str = Field("dp")
    precision: int = Field(16)
    override_float32_matmul: bool = Field(True)
    float32_matmul: str = Field("medium")


class RealFakeParams(Args):
    jsonl_file: Path
    dry_run: bool = Field(False)
    model_name: str = Field("convnext_tiny")
    batch_size: int = Field(256)
    freeze_epochs: int = Field(3)
    epochs: int = Field(6)
    base_lr: float = Field(1e-3)
    pretrained: bool = Field(True)
    progress_bar: bool = Field(False)
    accelerator: AcceleratorParams = Field(default_factory=AcceleratorParams)


class RealFakeDataModule(pl.LightningDataModule):

    def __init__(self, jsonl_records: Path, batch_size: int, num_workers: int = 0):
        super().__init__()
        self.jsonl_records = jsonl_records
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.dss = self.dls = None

    def setup(self, stage=None):
        records = [json.loads(line) for line in self.jsonl_records.open()]
        self.dss = get_dss(records)
        self.dls = get_dls(*self.dss, self.batch_size, self.num_workers)

    def train_dataloader(self):
        return self.dls[0]

    def val_dataloader(self):
        return self.dls[1]


class RealFakeClassifier(pl.LightningModule):

    def __init__(self, params: RealFakeParams):
        super().__init__()
        self.params = params
        self.ce = nn.BCEWithLogitsLoss()
        self.model = timm.create_model(params.model_name, pretrained=params.pretrained, num_classes=N_CLASSES)
        self.acc = torchmetrics.Accuracy(task="binary")

    def train_dataloader(self):
        return self.dls.train

    def val_dataloader(self):
        return self.dls.valid

    def forward(self, batch):
        x, y = batch["image"], batch["label"]
        y = torch.nn.functional.one_hot(y, num_classes=N_CLASSES).float()
        out = self.model(x)
        loss = self.ce(out, y)
        return loss, out, y
        
    def training_step(self, batch, batch_idx):
        loss, _, _ = self.forward(batch)
        self.log("train_loss", loss, on_epoch=True, on_step=False)
        return loss

    def validation_step(self, batch, batch_idx):
        loss, out, y = self.forward(batch)
        y_pred = out.sigmoid().argmax(dim=-1)
        y_true = y.argmax(dim=-1)
        self.log("val_loss", loss, on_epoch=True, on_step=False)
        return {"gt": y_true, "yhat": y_pred}

    def validation_step_end(self, outputs):
        self.acc.update(outputs["yhat"], outputs["gt"])

    def validation_epoch_end(self, outputs):
        self.log("val_acc", self.acc.compute(), on_epoch=True)
        self.acc.reset()

    def configure_optimizers(self):
        adamw = torch.optim.AdamW(self.parameters(), lr=self.params.base_lr)
        one_cycle = torch.optim.lr_scheduler.OneCycleLR(
            adamw, 
            max_lr=self.params.base_lr,
            total_steps=self.trainer.estimated_stepping_batches
        )
        return [adamw], [one_cycle]