"""Define the neural network model NetPP2D2_2_PRELU.

The NetPP2D2_2_PRELU class is a subclass of ChiNN and implements the forward pass of the neural network.
It consists of two fully connected layers with PReLU activation functions.


Methods:
    __init__(): Initializes the NetPP2D2_2_PRELU model.
    forward(x: torch.Tensor) -> torch.Tensor: Performs the forward pass of the neural network.
    init_weights(file: str) -> None: Initializes the weights of the model.
    print_param() -> None: Prints the parameters of the model.


"""

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal, cast

import torch
import yaml
from torch import nn

from coral.chi_nn import ChiNN
from coral.neural_networks.nn_model_type import (
    ActivationFunctionType,
    NNModelType,
    activation_map,
)
from coral.utils.logger import coral_logger

if TYPE_CHECKING:
    from collections.abc import Callable


@dataclass
class MultiLayerPerceptronArgs:
    """Arguments for the MultiLayerPerceptron model."""

    type: Literal[NNModelType.MULTI_LAYER_PERCEPTRON]
    number_neurons_per_layer: list[int]
    list_of_activation_functions: list[ActivationFunctionType]

    def __str__(self) -> str:
        """Return a readable summary of the model arguments."""
        neurons = "-".join(map(str, self.number_neurons_per_layer))
        activations = "-".join(act.value for act in self.list_of_activation_functions)
        return f"Type: {self.type}, Neurons: [{neurons}], Activations: [{activations}]"

    def filename(self) -> str:
        """Generate a filename for the model based on its architecture.

        Returns:
            str: A filename-safe string representing the model architecture.

        """
        # Create a filename-safe version by replacing spaces with underscores and ensuring everything is lowercase
        neurons = "_".join(map(str, self.number_neurons_per_layer))
        activations = "_".join(act.value for act in self.list_of_activation_functions)
        return f"{self.type}_{neurons}_{activations}"


def build_sequential(
    layers: list[int], activations: list[ActivationFunctionType]
) -> nn.Sequential:
    """Build a sequential neural network model.

    Args:
        layers (list[int]): A list of integers representing the number of neurons in each layer.
        activations (list[ActivationFunctionType]): A list of activation functions to use in each layer.

    Returns:
        nn.Sequential: A sequential neural network model.

    """
    modules: list[nn.Module] = []
    for i in range(len(layers) - 1):
        modules.append(nn.Linear(layers[i], layers[i + 1]))
        act: ActivationFunctionType = activations[i]
        if act in activation_map:
            modules.append(activation_map[act]())
    return nn.Sequential(*modules)


def extract_sequential_model_data(model: nn.Sequential) -> dict[str, Any]:
    """Extract the weights and biases from a Sequential model in a layered dictionary format.

    Args:
        model (nn.Sequential): The Sequential model to extract from.

    Returns:
        dict[str, Any]: A dictionary with layer names as keys and parameter data as values.

    """
    layers_data: dict[str, Any] = {}

    for idx, layer in enumerate(model):
        layer_info = {}
        if hasattr(layer, "weight"):  # and layer.weight is not None:
            weight_tensor = cast("torch.Tensor", layer.weight)
            layer_info["weight"] = weight_tensor.detach().cpu().numpy().tolist()
        if hasattr(layer, "bias"):  # and layer.bias is not None:
            bias_tensor = cast("torch.Tensor", layer.bias)
            layer_info["bias"] = bias_tensor.detach().cpu().numpy().tolist()

        layer_name = f"layer_{idx}_{layer.__class__.__name__}"
        layers_data[layer_name] = layer_info if layer_info else "No parameters"

    return layers_data


class MultiLayerPerceptron(ChiNN):
    """Neural network model for board evaluation using 2D2_2_PRELU architecture.

    Attributes:
        fc1 (nn.Linear): The first fully connected layer.
        relu_1 (nn.PReLU): The first PReLU activation function.
        fc2 (nn.Linear): The second fully connected layer.
        tanh (nn.Tanh): The tanh activation function.

    """

    def __init__(self, args: MultiLayerPerceptronArgs) -> None:
        """Construct the NetPP2D2_2_PRELU class and initialize the layers."""
        super().__init__()

        self.model: Callable[[torch.Tensor], torch.Tensor] = build_sequential(
            layers=args.number_neurons_per_layer,
            activations=args.list_of_activation_functions,
        )

    def __getstate__(self) -> dict[str, object]:
        """Get the state of the neural network for pickling."""
        return self.__dict__.copy()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of the neural network.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The output tensor.

        """
        return self.model(x)

    def init_weights(self) -> None:
        """Initialize the weights of the neural network."""
        return

    def log_readable_model_weights_to_file(self, file_path: str) -> None:
        """Write the model weights and biases into a YAML file.

        Args:
            file_path (str): The path where the YAML file will be saved.

        """
        assert isinstance(self.model, nn.Sequential)
        layers_data = extract_sequential_model_data(self.model)

        with open(file_path, "w", encoding="utf-8") as f:
            yaml.dump(layers_data, f, default_flow_style=False, sort_keys=False)

        coral_logger.info("Model weights successfully written to %s", file_path)

    def print_param(self) -> None:
        """Print the model weights and biases to the console in YAML format."""
        assert isinstance(self.model, nn.Sequential)
        layers_data = extract_sequential_model_data(self.model)
        yaml_str = yaml.dump(layers_data, default_flow_style=False, sort_keys=False)
        coral_logger.info(yaml_str)
