import os
from training.losses import disc_losses
from training.optimizers import optimizers
from pathlib import Path
from typing import Optional, List, Tuple, Dict
from dataclasses import dataclass, field
from omegaconf import OmegaConf, MISSING
from utils.class_registry import ClassRegistry
from models.methods import methods_registry
from metrics.metrics import metrics_registry


args = ClassRegistry()


@args.add_to_registry("exp")
@dataclass
class ExperimentArgs:
    config_dir: str = str(Path(__file__).resolve().parent / "configs")
    config: str = MISSING
    exp_dir: str = "experiments"
    name: str = MISSING
    seed: int = 1
    root: str = os.getenv("EXP_ROOT", ".")
    wandb: bool = True
    wandb_project: str = "sfe"
    domain: str = "human_faces"


@args.add_to_registry("data")
@dataclass
class DataArgs:
    special_dir: str = MISSING
    transform: str = "face_1024"
    input_train_dir: str = MISSING
    input_val_dir: str = MISSING


@args.add_to_registry("train")
@dataclass
class TrainingArgs:
    train_runner: str = "base_training_runner"
    encoder_optimizer: str = "ranger"
    disc_optimizer: str = "adam"
    resume_path: str = ""
    val_metrics: List[str] = field(
        default_factory=lambda: ["msssim", "lpips", "l2", "fid"]
    )
    start_step: int = 0
    steps: int = 300000
    log_step: int = 500
    checkpoint_step: int = 15000
    val_step: int = 15000
    train_dis: bool = False
    dis_train_start_step: int = 150000
    bs_used_before_adv_loss: int = 8
    disc_edits: List[str] = field(
        default_factory=lambda: []
    )

@args.add_to_registry("model")
@dataclass
class ModelArgs:
    method: str = "fse_full"
    device: str = "0"
    batch_size: int = 4
    workers: int = 4
    checkpoint_path: str = ""


@args.add_to_registry("encoder_losses")
@dataclass
class EncoderLossesArgs:
    l2: float = 0.0
    lpips: float = 0.0
    lpips_scale: float = 0.0
    id: float = 0.0
    moco: float = 0.0
    adv: float = 0.0
    feat_rec: float = 0.0
    feat_rec_l1: float = 0.0
    l2_latent: float = 0.0
    id_vit: float = 0.0


MethodsArgs = methods_registry.make_dataclass_from_args("MethodsArgs")
args.add_to_registry("methods_args")(MethodsArgs)

DiscLossesArgs = disc_losses.make_dataclass_from_args("DiscLossesArgs")
args.add_to_registry("disc_losses")(DiscLossesArgs)

OptimizersArgs = optimizers.make_dataclass_from_args("OptimizersArgs")
args.add_to_registry("optimizers")(OptimizersArgs)

MetricsArgs = metrics_registry.make_dataclass_from_args("MetricsArgs")
args.add_to_registry("metrics")(MetricsArgs)


Args = args.make_dataclass_from_classes("Args")


def load_config():
    config = OmegaConf.structured(Args)

    conf_cli = OmegaConf.from_cli()
    config.exp.config = conf_cli.exp.config
    config.exp.config_dir = conf_cli.exp.config_dir

    config_path = os.path.join(config.exp.config_dir, config.exp.config)
    conf_file = OmegaConf.load(config_path)
    config = OmegaConf.merge(config, conf_file)

    for method in list(config.methods_args.keys()):
        if method != config.model.method:
            config.methods_args.__delattr__(method)

    config = OmegaConf.merge(config, conf_cli)

    return config