from typing import Sequence, cast

import torch
from torch import nn

from .encoders import EncoderFactory
from .q_functions import QFunctionFactory, DropQFunctionFactory, DropContinuousMeanQFunction
from .torch import (
    CategoricalPolicy,
    ConditionalVAE,
    DeterministicPolicy,
    DeterministicRegressor,
    DeterministicResidualPolicy,
    DiscreteImitator,
    EnsembleContinuousQFunction,
    EnsembleDiscreteQFunction,
    NonSquashedNormalPolicy,
    Parameter,
    ProbabilisticDynamicsModel,
    ProbabilisticEnsembleDynamicsModel,
    ProbablisticRegressor,
    SquashedNormalPolicy,
    ValueFunction,
    DropDeterministicRegressor,
    DropProbablisticRegressor,
    DropDeterministicPolicy,
    DropSquashedNormalPolicy,
    DropDiscreteImitator,
    DropEnsembleContinuousQFunction, 
    DropEnergyFunction,
)



def create_discrete_q_function(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
    q_func_factory: QFunctionFactory,
    n_ensembles: int = 1,
) -> EnsembleDiscreteQFunction:
    if q_func_factory.share_encoder:
        encoder = encoder_factory.create(observation_shape)
        # normalize gradient scale by ensemble size
        for p in cast(nn.Module, encoder).parameters():
            p.register_hook(lambda grad: grad / n_ensembles)

    q_funcs = []
    for _ in range(n_ensembles):
        if not q_func_factory.share_encoder:
            encoder = encoder_factory.create(observation_shape)
        q_funcs.append(q_func_factory.create_discrete(encoder, action_size))
    return EnsembleDiscreteQFunction(q_funcs)


def create_continuous_q_function(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
    q_func_factory: QFunctionFactory,
    n_ensembles: int = 1,
) -> EnsembleContinuousQFunction:
    if q_func_factory.share_encoder:
        encoder = encoder_factory.create_with_action(
            observation_shape, action_size
        )
        # normalize gradient scale by ensemble size
        for p in cast(nn.Module, encoder).parameters():
            p.register_hook(lambda grad: grad / n_ensembles)

    q_funcs = []
    for _ in range(n_ensembles):
        if not q_func_factory.share_encoder:
            encoder = encoder_factory.create_with_action(
                observation_shape, action_size
            )
        q_funcs.append(q_func_factory.create_continuous(encoder))
    return EnsembleContinuousQFunction(q_funcs)

def create_drop_continuous_q_function_ens(
    observation_shape: Sequence[int],
    action_size: int,
    embedding_size: int,
    encoder_factory: EncoderFactory,
    q_func_factory: DropQFunctionFactory,
    n_ensembles: int = 1,
) -> DropEnsembleContinuousQFunction:
    if q_func_factory.share_encoder:
        encoder = encoder_factory.create_with_action(
            observation_shape, action_size
        )
        # normalize gradient scale by ensemble size
        for p in cast(nn.Module, encoder).parameters():
            p.register_hook(lambda grad: grad / n_ensembles)

    q_funcs = []
    for _ in range(n_ensembles):
        if not q_func_factory.share_encoder:
            encoder = encoder_factory.create_with_action(
                observation_shape, action_size
            )
        q_funcs.append(q_func_factory.create_continuous(encoder, embedding_size))
    return DropEnsembleContinuousQFunction(q_funcs)

def create_drop_continuous_q_function(
    observation_shape: Sequence[int],
    action_size: int,
    embedding_size: int,
    encoder_factory: EncoderFactory,
) -> DropContinuousMeanQFunction:
    encoder = encoder_factory.create_with_action(
                observation_shape, action_size
            )
    return DropContinuousMeanQFunction(encoder, embedding_size)

def create_drop_energy_function(
    observation_shape: Sequence[int],
    action_size: int,
    embedding_size: int,
    encoder_factory: EncoderFactory,
) -> DropEnergyFunction:
    encoder = encoder_factory.create_with_action_state(observation_shape, action_size)
    return DropEnergyFunction(encoder, embedding_size)

def create_deterministic_policy(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
) -> DeterministicPolicy:
    encoder = encoder_factory.create(observation_shape)
    return DeterministicPolicy(encoder, action_size)


def create_drop_deterministic_policy(
    observation_shape: Sequence[int],
    action_size: int,
    embedding_size: int,
    encoder_factory: EncoderFactory,
) -> DropDeterministicPolicy:
    encoder = encoder_factory.create(observation_shape)
    return DropDeterministicPolicy(encoder, action_size, embedding_size)


def create_deterministic_residual_policy(
    observation_shape: Sequence[int],
    action_size: int,
    scale: float,
    encoder_factory: EncoderFactory,
) -> DeterministicResidualPolicy:
    encoder = encoder_factory.create_with_action(observation_shape, action_size)
    return DeterministicResidualPolicy(encoder, scale)


def create_squashed_normal_policy(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
    min_logstd: float = -20.0,
    max_logstd: float = 2.0,
    use_std_parameter: bool = False,
) -> SquashedNormalPolicy:
    encoder = encoder_factory.create(observation_shape)
    return SquashedNormalPolicy(
        encoder,
        action_size,
        min_logstd=min_logstd,
        max_logstd=max_logstd,
        use_std_parameter=use_std_parameter,
    )

def create_drop_squashed_normal_policy(
    observation_shape: Sequence[int],
    action_size: int,
    embedding_size: int,
    encoder_factory: EncoderFactory,
    min_logstd: float = -20.0,
    max_logstd: float = 2.0,
    use_std_parameter: bool = False,
) -> DropSquashedNormalPolicy:
    encoder = encoder_factory.create(observation_shape)
    return DropSquashedNormalPolicy(
        encoder,
        action_size,
        embedding_size,
        min_logstd=min_logstd,
        max_logstd=max_logstd,
        use_std_parameter=use_std_parameter,
    )


def create_non_squashed_normal_policy(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
    min_logstd: float = -20.0,
    max_logstd: float = 2.0,
    use_std_parameter: bool = False,
) -> NonSquashedNormalPolicy:
    encoder = encoder_factory.create(observation_shape)
    return NonSquashedNormalPolicy(
        encoder,
        action_size,
        min_logstd=min_logstd,
        max_logstd=max_logstd,
        use_std_parameter=use_std_parameter,
    )


def create_categorical_policy(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
) -> CategoricalPolicy:
    encoder = encoder_factory.create(observation_shape)
    return CategoricalPolicy(encoder, action_size)


def create_conditional_vae(
    observation_shape: Sequence[int],
    action_size: int,
    latent_size: int,
    beta: float,
    encoder_factory: EncoderFactory,
    min_logstd: float = -20.0,
    max_logstd: float = 2.0,
) -> ConditionalVAE:
    encoder_encoder = encoder_factory.create_with_action(
        observation_shape, action_size
    )
    decoder_encoder = encoder_factory.create_with_action(
        observation_shape, latent_size
    )
    return ConditionalVAE(
        encoder_encoder,
        decoder_encoder,
        beta,
        min_logstd=min_logstd,
        max_logstd=max_logstd,
    )


def create_discrete_imitator(
    observation_shape: Sequence[int],
    action_size: int,
    beta: float,
    encoder_factory: EncoderFactory,
) -> DiscreteImitator:
    encoder = encoder_factory.create(observation_shape)
    return DiscreteImitator(encoder, action_size, beta)

def create_drop_discrete_imitator(
    observation_shape: Sequence[int],
    action_size: int,
    embedding_size: int,
    beta: float,
    encoder_factory: EncoderFactory,
) -> DropDiscreteImitator:
    encoder = encoder_factory.create(observation_shape)
    return DropDiscreteImitator(encoder, action_size, embedding_size, beta)


def create_deterministic_regressor(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
) -> DeterministicRegressor:
    encoder = encoder_factory.create(observation_shape)
    return DeterministicRegressor(encoder, action_size)


def create_probablistic_regressor(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
    min_logstd: float = -20.0,
    max_logstd: float = 2.0,
) -> ProbablisticRegressor:
    encoder = encoder_factory.create(observation_shape)
    return ProbablisticRegressor(
        encoder, action_size, min_logstd=min_logstd, max_logstd=max_logstd
    )

def create_drop_deterministic_regressor(
    observation_shape: Sequence[int],
    action_size: int,
    embedding_size: int,
    encoder_factory: EncoderFactory,
) -> DropDeterministicRegressor:
    encoder = encoder_factory.create(observation_shape)
    return DropDeterministicRegressor(encoder, action_size, embedding_size)


def create_drop_probablistic_regressor(
    observation_shape: Sequence[int],
    action_size: int,
    embedding_size: int,
    encoder_factory: EncoderFactory,
    min_logstd: float = -20.0,
    max_logstd: float = 2.0,
) -> DropProbablisticRegressor:
    encoder = encoder_factory.create(observation_shape)
    return DropProbablisticRegressor(
        encoder, action_size, embedding_size, min_logstd=min_logstd, max_logstd=max_logstd
    )


def create_value_function(
    observation_shape: Sequence[int], encoder_factory: EncoderFactory
) -> ValueFunction:
    encoder = encoder_factory.create(observation_shape)
    return ValueFunction(encoder)


def create_probabilistic_ensemble_dynamics_model(
    observation_shape: Sequence[int],
    action_size: int,
    encoder_factory: EncoderFactory,
    n_ensembles: int = 5,
    discrete_action: bool = False,
) -> ProbabilisticEnsembleDynamicsModel:
    models = []
    for _ in range(n_ensembles):
        encoder = encoder_factory.create_with_action(
            observation_shape=observation_shape,
            action_size=action_size,
            discrete_action=discrete_action,
        )
        model = ProbabilisticDynamicsModel(encoder)
        models.append(model)
    return ProbabilisticEnsembleDynamicsModel(models)


def create_parameter(shape: Sequence[int], initial_value: float) -> Parameter:
    data = torch.full(shape, initial_value, dtype=torch.float32)
    return Parameter(data)
