from dataclasses import dataclass, field

from typing import Dict, Any, Optional, Type
from torch import optim, nn
import torchmetrics
import sklearn

from .embedding_utils import BaseEmbeddingConfig
from .backbone_utils import BaseBackboneConfig

import sys
if sys.version_info >= (3, 8):
    from typing import Literal
else:
    from typing_extensions import Literal


@dataclass
class BaseConfig:
    """ Configuration class for initializing components of the TabularS3L Lightning Module, including optimizers, 
    learning rate schedulers, and loss functions, along with their respective hyperparameters.

    Attributes:
        task (str): Specify whether the problem is regression or classification.
        embedding_config (BaseEmbeddingConfig): Configuration for the embedding layer.
        backbone_config (BaseBackboneConfig): Configuration for the backbone network.
        output_dim (int): The dimension of output.
        loss_fn (str): Name of the loss function to be used. Must be an attribute of 'torch.nn'.
        loss_hparams (Dict[str, Any]): Hyperparameters for the loss function. Default is empty dictionary.
        metric (str): Name of the metric to be used. Must be an attribute of 'torchmetrics.functional' or 'sklearn.metrics'. Default is None.
        metric_hparams (Dict[str, Any]): Hyperparameters for the metric. Default is an empty dictionary.
        optim (str): Name of the optimizer to be used. Must be an attribute of 'torch.optim'. Default is 'AdamW'.
        optim_hparams (Dict[str, Any]): Hyperparameters for the optimizer. Default is {'lr': 0.0001, 'weight_decay': 0.00005}.
        scheduler (str): Name of the learning rate scheduler to be used. Must be an attribute of 'torch.optim.lr_scheduler' or None. Default is None.
        scheduler_hparams (Dict[str, Any]): Hyperparameters for the scheduler. Default is None, indicating no scheduler is used.
        initialization (str): The way to initialize neural network parameters. Default is 'kaiming_uniform'.
        random_seed (int): Seed for random number generators to ensure reproducibility. Defaults to 42.
    Raises:
        ValueError: If the specified 'optim' is not a valid optimizer in 'torch.optim'.
        ValueError: If the specified 'scheduler' is not None and is not a valid scheduler in 'torch.optim.lr_scheduler'.
        ValueError: If the specified 'loss_fn' is not None and is not a valid loss function in 'torch.nn'.        
        ValueError: If the specified 'metric' is not a valid metric in 'torchmetrics' or 'sklearn.metrics'.
        ValueError: If the specified 'task' is not a valid task in ['regression', 'classification']'.
        ValueError: If the combination of `backbone_config` and `embedding_config` is invalid.
    """
    task: str

    embedding_config: BaseEmbeddingConfig

    backbone_config: BaseBackboneConfig

    output_dim: int

    loss_fn: str

    loss_hparams: Dict[str, Any] = field(default_factory=dict)

    metric: Optional[str] = field(default=None)

    metric_hparams: Dict[str, Any] = field(default_factory=dict)

    optim: str = field(default="AdamW")

    optim_hparams: Dict[str, Any] = field(
        default_factory=lambda: {
            "lr": 0.0001,
            "weight_decay": 0.00005
        }
    )

    scheduler: Optional[str] = field(default=None)

    scheduler_hparams: Optional[Dict[str, Any]] = field(default=None)

    initialization: Literal['xavier_uniform', 'xavier_normal',
                            'kaiming_uniform', 'kaiming_normal', 'uniform',
                            'normal'] = "kaiming_uniform"

    random_seed: int = field(default=42)

    def __post_init__(self):

        if (type(self.task) is not str or (self.task != "regression" and self.task != "classification")):
            raise ValueError(
                f"{self.task} is not a valid task. Choices are: ['regression', 'classification']")

        if type(self.optim) is not str or not hasattr(optim, self.optim):
            raise ValueError(
                f"{self.optim} is not a valid optimizer in torch.optim")

        if self.scheduler is not None and (
                type(self.scheduler)
                is not str or not hasattr(optim.lr_scheduler, self.scheduler)):
            raise ValueError(
                f"{self.scheduler} is not a valid scheduler in torch.optim.lr_scheduler")

        if type(self.loss_fn) is not str or not hasattr(nn, self.loss_fn):
            raise ValueError(
                f"{self.loss_fn} is not a valid loss function in torch.nn")

        if self.metric is None:
            if self.task == "regression":
                self.metric = "mean_squared_error"
            else:
                self.metric = "accuracy_score"

        elif (type(self.metric) is not str or (not hasattr(torchmetrics.functional, self.metric) and not hasattr(sklearn.metrics, self.metric))):
            raise ValueError(
                f"{self.metric} is not a valid metric in torchmetrics.functional or sklearn.metrics")

        if self.backbone_config.name == "transformer" and self.embedding_config.name == "identity":
            raise ValueError(
                f"Transformer backbone and Identity Embedding are not compatible.")

        if self.backbone_config.name == "transformer" and self.embedding_config.required_token_dim == 1:
            raise ValueError(
                f"Embedding's 'required_token_dim' should be 2 for transformer backbone, not {self.embedding_config.required_token_dim}")
