# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
All the functions to build the relevant solvers and used objects
from the Hydra config.
"""

from enum import Enum
import logging
import typing as tp

import dora
import flashy
import omegaconf
import torch
from torch import nn
from torch.optim import Optimizer
# LRScheduler was renamed in some torch versions
try:
    from torch.optim.lr_scheduler import LRScheduler  # type: ignore
except ImportError:
    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler

from .base import StandardSolver
from .. import adversarial, data, losses, metrics, optim
from ..utils.utils import dict_from_config, get_loader


logger = logging.getLogger(__name__)


class DatasetType(Enum):
    AUDIO = "audio"
    MUSIC = "music"
    SOUND = "sound"


def get_solver(cfg: omegaconf.DictConfig) -> StandardSolver:
    """Instantiate solver from config."""
    from .audiogen import AudioGenSolver
    from .compression import CompressionSolver
    from .musicgen import MusicGenSolver
    from .diffusion import DiffusionSolver
    klass = {
        'compression': CompressionSolver,
        'musicgen': MusicGenSolver,
        'audiogen': AudioGenSolver,
        'lm': MusicGenSolver,  # backward compatibility
        'diffusion': DiffusionSolver,
        'sound_lm': AudioGenSolver,  # backward compatibility
    }[cfg.solver]
    return klass(cfg)  # type: ignore


def get_optim_parameter_groups(model: nn.Module):
    """Create parameter groups for the model using the appropriate method
    if defined for each modules, to create the different groups.

    Args:
        model (nn.Module): torch model
    Returns:
        List of parameter groups
    """
    seen_params: tp.Set[nn.parameter.Parameter] = set()
    other_params = []
    groups = []
    for name, module in model.named_modules():
        if hasattr(module, 'make_optim_group'):
            group = module.make_optim_group()
            params = set(group['params'])
            assert params.isdisjoint(seen_params)
            seen_params |= set(params)
            groups.append(group)
    for param in model.parameters():
        if param not in seen_params:
            other_params.append(param)
    groups.insert(0, {'params': other_params})
    parameters = groups
    return parameters


def get_optimizer(params: tp.Union[nn.Module, tp.Iterable[torch.Tensor]], cfg: omegaconf.DictConfig) -> Optimizer:
    """Build torch optimizer from config and set of parameters.
    Supported optimizers: Adam, AdamW

    Args:
        params (nn.Module or iterable of torch.Tensor): Parameters to optimize.
        cfg (DictConfig): Optimization-related configuration.
    Returns:
        torch.optim.Optimizer.
    """
    if 'optimizer' not in cfg:
        if getattr(cfg, 'optim', None) is not None:
            raise KeyError("Optimizer not found in config. Try instantiating optimizer from cfg.optim?")
        else:
            raise KeyError("Optimizer not found in config.")

    parameters = get_optim_parameter_groups(params) if isinstance(params, nn.Module) else params
    optimizer: torch.optim.Optimizer
    if cfg.optimizer == 'adam':
        optimizer = torch.optim.Adam(parameters, lr=cfg.lr, **cfg.adam)
    elif cfg.optimizer == 'adamw':
        optimizer = torch.optim.AdamW(parameters, lr=cfg.lr, **cfg.adam)
    elif cfg.optimizer == 'dadam':
        optimizer = optim.DAdaptAdam(parameters, lr=cfg.lr, **cfg.adam)
    else:
        raise ValueError(f"Unsupported LR Scheduler: {cfg.lr_scheduler}")
    return optimizer


def get_lr_scheduler(optimizer: torch.optim.Optimizer,
                     cfg: omegaconf.DictConfig,
                     total_updates: int) -> tp.Optional[LRScheduler]:
    """Build torch learning rate scheduler from config and associated optimizer.
    Supported learning rate schedulers: ExponentialLRScheduler, PlateauLRScheduler

    Args:
        optimizer (torch.optim.Optimizer): Optimizer.
        cfg (DictConfig): Schedule-related configuration.
        total_updates (int): Total number of updates.
    Returns:
        torch.optim.Optimizer.
    """
    if 'lr_scheduler' not in cfg:
        raise KeyError("LR Scheduler not found in config")

    lr_sched: tp.Optional[LRScheduler] = None
    if cfg.lr_scheduler == 'step':
        lr_sched = torch.optim.lr_scheduler.StepLR(optimizer, **cfg.step)
    elif cfg.lr_scheduler == 'exponential':
        lr_sched = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=cfg.exponential)
    elif cfg.lr_scheduler == 'cosine':
        kwargs = dict_from_config(cfg.cosine)
        warmup_steps = kwargs.pop('warmup')
        lr_sched = optim.CosineLRScheduler(
            optimizer, warmup_steps=warmup_steps, total_steps=total_updates, **kwargs)
    elif cfg.lr_scheduler == 'polynomial_decay':
        kwargs = dict_from_config(cfg.polynomial_decay)
        warmup_steps = kwargs.pop('warmup')
        lr_sched = optim.PolynomialDecayLRScheduler(
            optimizer, warmup_steps=warmup_steps, total_steps=total_updates, **kwargs)
    elif cfg.lr_scheduler == 'inverse_sqrt':
        kwargs = dict_from_config(cfg.inverse_sqrt)
        warmup_steps = kwargs.pop('warmup')
        lr_sched = optim.InverseSquareRootLRScheduler(optimizer, warmup_steps=warmup_steps, **kwargs)
    elif cfg.lr_scheduler == 'linear_warmup':
        kwargs = dict_from_config(cfg.linear_warmup)
        warmup_steps = kwargs.pop('warmup')
        lr_sched = optim.LinearWarmupLRScheduler(optimizer, warmup_steps=warmup_steps, **kwargs)
    elif cfg.lr_scheduler is not None:
        raise ValueError(f"Unsupported LR Scheduler: {cfg.lr_scheduler}")
    return lr_sched


def get_ema(module_dict: nn.ModuleDict, cfg: omegaconf.DictConfig) -> tp.Optional[optim.ModuleDictEMA]:
    """Initialize Exponential Moving Average.

    Args:
        module_dict (nn.ModuleDict): ModuleDict for which to compute the EMA.
        cfg (omegaconf.DictConfig): Optim EMA configuration.
    Returns:
        optim.ModuleDictEMA: EMA version of the ModuleDict.
    """
    kw: tp.Dict[str, tp.Any] = dict(cfg)
    use = kw.pop('use', False)
    decay = kw.pop('decay', None)
    device = kw.pop('device', None)
    if not use:
        return None
    if len(module_dict) == 0:
        raise ValueError("Trying to build EMA but an empty module_dict source is provided!")
    ema_module = optim.ModuleDictEMA(module_dict, decay=decay, device=device)
    return ema_module


def get_loss(loss_name: str, cfg: omegaconf.DictConfig):
    """Instantiate loss from configuration."""
    klass = {
        'l1': torch.nn.L1Loss,
        'l2': torch.nn.MSELoss,
        'mel': losses.MelSpectrogramL1Loss,
        'mrstft': losses.MRSTFTLoss,
        'msspec': losses.MultiScaleMelSpectrogramLoss,
        'sisnr': losses.SISNR,
    }[loss_name]
    kwargs = dict(getattr(cfg, loss_name))
    return klass(**kwargs)


def get_balancer(loss_weights: tp.Dict[str, float], cfg: omegaconf.DictConfig) -> losses.Balancer:
    """Instantiate loss balancer from configuration for the provided weights."""
    kwargs: tp.Dict[str, tp.Any] = dict_from_config(cfg)
    return losses.Balancer(loss_weights, **kwargs)


def get_adversary(name: str, cfg: omegaconf.DictConfig) -> nn.Module:
    """Initialize adversary from config."""
    klass = {
        'msd': adversarial.MultiScaleDiscriminator,
        'mpd': adversarial.MultiPeriodDiscriminator,
        'msstftd': adversarial.MultiScaleSTFTDiscriminator,
    }[name]
    adv_cfg: tp.Dict[str, tp.Any] = dict(getattr(cfg, name))
    return klass(**adv_cfg)


def get_adversarial_losses(cfg) -> nn.ModuleDict:
    """Initialize dict of adversarial losses from config."""
    device = cfg.device
    adv_cfg = getattr(cfg, 'adversarial')
    adversaries = adv_cfg.get('adversaries', [])
    adv_loss_name = adv_cfg['adv_loss']
    feat_loss_name = adv_cfg.get('feat_loss')
    normalize = adv_cfg.get('normalize', True)
    feat_loss: tp.Optional[adversarial.FeatureMatchingLoss] = None
    if feat_loss_name:
        assert feat_loss_name in ['l1', 'l2'], f"Feature loss only support L1 or L2 but {feat_loss_name} found."
        loss = get_loss(feat_loss_name, cfg)
        feat_loss = adversarial.FeatureMatchingLoss(loss, normalize)
    loss = adversarial.get_adv_criterion(adv_loss_name)
    loss_real = adversarial.get_real_criterion(adv_loss_name)
    loss_fake = adversarial.get_fake_criterion(adv_loss_name)
    adv_losses = nn.ModuleDict()
    for adv_name in adversaries:
        adversary = get_adversary(adv_name, cfg).to(device)
        optimizer = get_optimizer(adversary.parameters(), cfg.optim)
        adv_loss = adversarial.AdversarialLoss(
            adversary,
            optimizer,
            loss=loss,
            loss_real=loss_real,
            loss_fake=loss_fake,
            loss_feat=feat_loss,
            normalize=normalize
        )
        adv_losses[adv_name] = adv_loss
    return adv_losses


def get_visqol(cfg: omegaconf.DictConfig) -> metrics.ViSQOL:
    """Instantiate ViSQOL metric from config."""
    kwargs = dict_from_config(cfg)
    return metrics.ViSQOL(**kwargs)


def get_fad(cfg: omegaconf.DictConfig) -> metrics.FrechetAudioDistanceMetric:
    """Instantiate Frechet Audio Distance metric from config."""
    kwargs = dict_from_config(cfg.tf)
    xp = dora.get_xp()
    kwargs['log_folder'] = xp.folder
    return metrics.FrechetAudioDistanceMetric(**kwargs)


def get_kldiv(cfg: omegaconf.DictConfig) -> metrics.KLDivergenceMetric:
    """Instantiate KL-Divergence metric from config."""
    kld_metrics = {
        'passt': metrics.PasstKLDivergenceMetric,
    }
    klass = kld_metrics[cfg.model]
    kwargs = dict_from_config(cfg.get(cfg.model))
    return klass(**kwargs)


def get_text_consistency(cfg: omegaconf.DictConfig) -> metrics.TextConsistencyMetric:
    """Instantiate Text Consistency metric from config."""
    text_consistency_metrics = {
        'clap': metrics.CLAPTextConsistencyMetric
    }
    klass = text_consistency_metrics[cfg.model]
    kwargs = dict_from_config(cfg.get(cfg.model))
    return klass(**kwargs)


def get_chroma_cosine_similarity(cfg: omegaconf.DictConfig) -> metrics.ChromaCosineSimilarityMetric:
    """Instantiate Chroma Cosine Similarity metric from config."""
    assert cfg.model == 'chroma_base', "Only support 'chroma_base' method for chroma cosine similarity metric"
    kwargs = dict_from_config(cfg.get(cfg.model))
    return metrics.ChromaCosineSimilarityMetric(**kwargs)


def get_audio_datasets(cfg: omegaconf.DictConfig,
                       dataset_type: DatasetType = DatasetType.AUDIO) -> tp.Dict[str, torch.utils.data.DataLoader]:
    """Build AudioDataset from configuration.

    Args:
        cfg (omegaconf.DictConfig): Configuration.
        dataset_type: The type of dataset to create.
    Returns:
        dict[str, torch.utils.data.DataLoader]: Map of dataloader for each data split.
    """
    dataloaders: dict = {}

    sample_rate = cfg.sample_rate
    channels = cfg.channels
    seed = cfg.seed
    max_sample_rate = cfg.datasource.max_sample_rate
    max_channels = cfg.datasource.max_channels

    assert cfg.dataset is not None, "Could not find dataset definition in config"

    dataset_cfg = dict_from_config(cfg.dataset)
    splits_cfg: dict = {}
    splits_cfg['train'] = dataset_cfg.pop('train')
    splits_cfg['valid'] = dataset_cfg.pop('valid')
    splits_cfg['evaluate'] = dataset_cfg.pop('evaluate')
    splits_cfg['generate'] = dataset_cfg.pop('generate')
    execute_only_stage = cfg.get('execute_only', None)

    for split, path in cfg.datasource.items():
        if not isinstance(path, str):
            continue  # skipping this as not a path
        if execute_only_stage is not None and split != execute_only_stage:
            continue
        logger.info(f"Loading audio data split {split}: {str(path)}")
        assert (
            cfg.sample_rate <= max_sample_rate
        ), f"Expecting a max sample rate of {max_sample_rate} for datasource but {sample_rate} found."
        assert (
            cfg.channels <= max_channels
        ), f"Expecting a max number of channels of {max_channels} for datasource but {channels} found."

        split_cfg = splits_cfg[split]
        split_kwargs = {k: v for k, v in split_cfg.items()}
        kwargs = {**dataset_cfg, **split_kwargs}  # split kwargs overrides default dataset_cfg
        kwargs['sample_rate'] = sample_rate
        kwargs['channels'] = channels

        if kwargs.get('permutation_on_files') and cfg.optim.updates_per_epoch:
            kwargs['num_samples'] = (
                flashy.distrib.world_size() * cfg.dataset.batch_size * cfg.optim.updates_per_epoch)

        num_samples = kwargs['num_samples']
        shuffle = kwargs['shuffle']

        return_info = kwargs.pop('return_info')
        batch_size = kwargs.pop('batch_size', None)
        num_workers = kwargs.pop('num_workers')

        if dataset_type == DatasetType.MUSIC:
            dataset = data.music_dataset.MusicDataset.from_meta(path, **kwargs)
        elif dataset_type == DatasetType.SOUND:
            dataset = data.sound_dataset.SoundDataset.from_meta(path, **kwargs)
        elif dataset_type == DatasetType.AUDIO:
            dataset = data.info_audio_dataset.InfoAudioDataset.from_meta(path, return_info=return_info, **kwargs)
        else:
            raise ValueError(f"Dataset type is unsupported: {dataset_type}")

        loader = get_loader(
            dataset,
            num_samples,
            batch_size=batch_size,
            num_workers=num_workers,
            seed=seed,
            collate_fn=dataset.collater if return_info else None,
            shuffle=shuffle,
        )
        dataloaders[split] = loader

    return dataloaders