import math
from pathlib import Path
from typing import Any, Callable, Dict, Optional, Union

import numpy as np
import pyro
import torch
from pyro import distributions as pdist

import sbibm
from sbibm.tasks.simulator import Simulator
from sbibm.tasks.task import Task
from sbibm.utils.pyro import make_log_prob_grad_fn


class TwoMoons(Task):
    def __init__(self):
        """Two Moons
        """

        # Observation seeds to use when generating ground truth
        observation_seeds = [
            1000011,  # observation 1
            1000001,  # observation 2
            1000002,  # observation 3
            1000003,  # observation 4
            1000013,  # observation 5
            1000005,  # observation 6
            1000006,  # observation 7
            1000007,  # observation 8
            1000008,  # observation 9
            1000009,  # observation 10
        ]

        super().__init__(
            dim_parameters=2,
            dim_data=2,
            name=Path(__file__).parent.name,
            name_display="Two Moons",
            num_observations=10,
            num_posterior_samples=10000,
            num_reference_posterior_samples=10000,
            num_simulations=[100, 1000, 10000, 100000, 1000000],
            observation_seeds=observation_seeds,
            path=Path(__file__).parent.absolute(),
        )

        prior_bound = 1.0
        self.prior_params = {
            "low": -prior_bound * torch.ones((self.dim_parameters,)),
            "high": +prior_bound * torch.ones((self.dim_parameters,)),
        }
        self.prior_dist = pdist.Uniform(**self.prior_params).to_event(1)

        self.simulator_params = {
            "a_low": -math.pi / 2.0,
            "a_high": +math.pi / 2.0,
            "base_offset": 0.25,
            "r_loc": 0.1,
            "r_scale": 0.01,
        }

    def get_prior(self) -> Callable:
        def prior(num_samples=1):
            return pyro.sample("parameters", self.prior_dist.expand_by([num_samples]))

        return prior

    def get_simulator(self, max_calls: Optional[int] = None) -> Simulator:
        """Get function returning samples from simulator given parameters

        Args:
            max_calls: Maximum number of function calls. Additional calls will 
                result in SimulationBudgetExceeded exceptions. Defaults to None 
                for infinite budget

        Return:
            Simulator callable
        """

        def simulator(parameters):
            num_samples = parameters.shape[0]

            a_dist = (
                pdist.Uniform(
                    low=self.simulator_params["a_low"],
                    high=self.simulator_params["a_high"],
                )
                .expand_by((num_samples, 1))
                .to_event(1)
            )
            a = a_dist.sample()

            r_dist = (
                pdist.Normal(
                    self.simulator_params["r_loc"], self.simulator_params["r_scale"]
                )
                .expand_by((num_samples, 1))
                .to_event(1)
            )
            r = r_dist.sample()

            p = torch.cat(
                (
                    torch.cos(a) * r + self.simulator_params["base_offset"],
                    torch.sin(a) * r,
                ),
                dim=1,
            )

            return self._map_fun(parameters, p)

        return Simulator(task=self, simulator=simulator, max_calls=max_calls)

    @staticmethod
    def _map_fun(parameters: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
        ang = torch.tensor([-math.pi / 4.0])
        c = torch.cos(ang)
        s = torch.sin(ang)
        z0 = (c * parameters[:, 0] - s * parameters[:, 1]).reshape(-1, 1)
        z1 = (s * parameters[:, 0] + c * parameters[:, 1]).reshape(-1, 1)
        return p + torch.cat((-torch.abs(z0), z1), dim=1)

    @staticmethod
    def _map_fun_inv(parameters: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        ang = torch.tensor([-math.pi / 4.0])
        c = torch.cos(ang)
        s = torch.sin(ang)
        z0 = (c * parameters[:, 0] - s * parameters[:, 1]).reshape(-1, 1)
        z1 = (s * parameters[:, 0] + c * parameters[:, 1]).reshape(-1, 1)
        return x - torch.cat((-torch.abs(z0), z1), dim=1)

    def _likelihood(
        self, parameters: torch.Tensor, data: torch.Tensor, log: bool = True,
    ) -> torch.Tensor:
        if parameters.ndim == 1:
            parameters = parameters.reshape(1, -1)

        assert parameters.shape[1] == self.dim_parameters
        assert data.shape[1] == self.dim_data

        p = self._map_fun_inv(parameters, data).squeeze(0)
        if p.ndim == 1:
            p = p.reshape(1, -1)
        u = p[:, 0] - self.simulator_params["base_offset"]
        v = p[:, 1]

        r = torch.sqrt(u ** 2 + v ** 2)
        L = -0.5 * (
            (r - self.simulator_params["r_loc"]) / self.simulator_params["r_scale"]
        ) ** 2 - 0.5 * torch.log(
            2 * torch.tensor([math.pi]) * self.simulator_params["r_scale"] ** 2
        )

        if len(torch.where(u < 0.0)[0]) > 0:
            L[torch.where(u < 0.0)[0]] = -torch.tensor(math.inf) if log else 0.0

        return L if log else torch.exp(L)

    def _get_transforms(self, *args, **kwargs: Any,) -> Dict[str, Any]:
        return {"parameters": torch.distributions.transforms.identity_transform}

    def _get_log_prob_fn(
        self,
        num_observation: Optional[int] = None,
        observation: Optional[torch.Tensor] = None,
        **kwargs: Any,
    ) -> Callable:
        """Get potential function and initial parameters

        The potential function returns the unnormalized negative log
        posterior probability, and is useful to establish and verify
        the reference posterior. 
        
        Args:
            num_observation: Observation number
            observation: Instead of passing an observation number, an observation may be
                passed directly

        Returns:
            Potential function and proposal for initial parameters, e.g., to start MCMC
        """
        assert not (num_observation is None and observation is None)
        assert not (num_observation is not None and observation is not None)

        prior_dist = self.get_prior_dist()

        if num_observation is not None:
            observation = self.get_observation(num_observation=num_observation)

        observation = self.unflatten_data(observation)

        def log_prob_fn(parameters):
            if type(parameters) == dict:
                parameters = parameters["parameters"]
            return self._likelihood(
                parameters=parameters, data=observation, log=True
            ) + prior_dist.log_prob(parameters)

        return log_prob_fn

    def _get_log_prob_grad_fn(
        self,
        num_observation: Optional[int] = None,
        observation: Optional[torch.Tensor] = None,
        **kwargs: Any,
    ) -> Callable:
        lpgf = make_log_prob_grad_fn(
            self._get_log_prob_fn(
                num_observation=num_observation, observation=observation, **kwargs
            )
        )

        def log_prob_grad_fn(parameters):
            num_params = parameters.shape[0]
            grads = []
            for i in range(num_params):
                _, grad = lpgf({"parameters": parameters[i]})
                grads.append(grad)
            if len(grads) > 1:
                return torch.cat(grads).reshape(1, -1)
            else:
                return grad

        return log_prob_grad_fn

    def _sample_reference_posterior(
        self,
        num_samples: int,
        num_observation: int,
        observation: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Sample reference posterior for given observation

        Uses closed form solution

        Args:
            num_samples: Number of samples to generate
            num_observation: Observation number
            observation: Observed data, if None, will be loaded using `num_observation`
        
        Returns:
            Samples from reference posterior
        """
        log = sbibm.get_logger(__name__)

        if observation is None:
            observation = self.get_observation(num_observation)

        reference_posterior_samples = []

        ang = torch.tensor([-math.pi / 4.0])
        c = torch.cos(-ang)
        s = torch.sin(-ang)

        simulator = self.get_simulator()

        reference_posterior_samples = []
        counter = 0
        while len(reference_posterior_samples) < num_samples:
            counter += 1

            p = simulator(torch.zeros(1, 2))
            q = torch.zeros(2)
            q[0] = p[0, 0] - observation[0, 0]
            q[1] = observation[0, 1] - p[0, 1]

            if np.random.rand() < 0.5:
                q[0] = -q[0]

            sample = torch.tensor([[c * q[0] - s * q[1], s * q[0] + c * q[1]]])

            is_outside_prior = torch.isinf(self.prior_dist.log_prob(sample).sum())

            if len(reference_posterior_samples) > 0:
                is_duplicate = sample in torch.cat(reference_posterior_samples)
            else:
                is_duplicate = False

            if not is_outside_prior and not is_duplicate:
                reference_posterior_samples.append(sample)

        reference_posterior_samples = torch.cat(reference_posterior_samples)
        acceptance_rate = float(num_samples / counter)

        log.info(
            f"Acceptance rate for observation {num_observation}: {acceptance_rate}"
        )

        return reference_posterior_samples


if __name__ == "__main__":
    task = TwoMoons()
    task._setup()
