import os
import sys
import unittest
import torch
import scipy.stats as scst
import numpy as np

PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)

from pytorch_wrapper.ptd_binomial import PtdBinomial
from pytorch_wrapper.ptdm_pdf import PtdmPdf
from torch.distributions import (
    Binomial, Categorical, Poisson, NegativeBinomial, Multinomial, 
    Bernoulli, Beta, Gamma, Normal, Laplace, Exponential, LogNormal, 
    Pareto, Weibull, Dirichlet, VonMises, StudentT, Uniform
)


class TestPtdfPdf(unittest.TestCase):
    def setUp(self):
        """Set up test instance."""
        self.test_cases = {
            "Beta": [(2.0, 5.0, 0.3, scst.beta.pdf(0.3, 2, 5))],
            "Gamma": [(2.0, 1.0, 3.0, scst.gamma.pdf(3.0, 2, scale=1))],
            "Normal": [(0.0, 1.0, -1.5, scst.norm.pdf(-1.5, 0, 1))],
            "Exponential": [(1.0, 3.0, scst.expon.pdf(3.0, scale=1))],
            "Dirichlet": [(torch.tensor([1.0, 1.0, 1.0]), torch.tensor([0.2, 0.5, 0.3]), scst.dirichlet.pdf([0.2, 0.5, 0.3], [1.0, 1.0, 1.0]))],
            "VonMises": [(0.0, 1.0, 1.0, scst.vonmises.pdf(1.0, kappa=1.0, loc=0))],
            "StudentT": [(1.0, 0.0, scst.t.pdf(0.0, 1))],
            "Uniform": [(0.0, 1.0, 0.5, scst.uniform.pdf(0.5, loc=0, scale=1))],
        }

    def test_distributions(self):
        """Test probability mass and density functions for all PyTorch distributions."""
        for dist_name, cases in self.test_cases.items():  # For each distribution
            for case in cases:  # For each distribution's test cases
                with self.subTest(distribution=dist_name, params=case):
                    if dist_name == "Binomial":
                        distribution = PtdBinomial().f(str(case[0]), str(case[2]), "")[0]  # n, p
                        pdf = PtdmPdf().f(distribution, str(case[1]))[0]  # k
                    else:
                        # Instantiate PyTorch distribution directly if not in wrapper
                        if dist_name == "Poisson":
                            distribution = Poisson(torch.tensor(case[0]))
                            k_tensor = torch.tensor(case[1], dtype=torch.int64)
                        elif dist_name == "Categorical":
                            distribution = Categorical(case[0])
                            k_tensor = torch.tensor(case[1], dtype=torch.int64)
                        elif dist_name == "NegativeBinomial":
                            distribution = NegativeBinomial(torch.tensor(case[0]), torch.tensor(case[2]))
                            k_tensor = torch.tensor(case[1], dtype=torch.int64)
                        elif dist_name == "Bernoulli":
                            distribution = Bernoulli(torch.tensor(case[0]))
                            k_tensor = torch.tensor(case[1], dtype=torch.float32)
                        elif dist_name == "Beta":
                            distribution = Beta(torch.tensor(case[0]), torch.tensor(case[1]))
                            k_tensor = torch.tensor(case[2], dtype=torch.float32)
                        elif dist_name == "Gamma":
                            distribution = Gamma(torch.tensor(case[0]), torch.tensor(case[1]))
                            k_tensor = torch.tensor(case[2], dtype=torch.float32)
                        elif dist_name == "Normal":
                            distribution = Normal(torch.tensor(case[0]), torch.tensor(case[1]))
                            k_tensor = torch.tensor(case[2], dtype=torch.float32)
                        elif dist_name == "Laplace":
                            distribution = Laplace(torch.tensor(case[0]), torch.tensor(case[1]))
                            k_tensor = torch.tensor(case[2], dtype=torch.float32)
                        elif dist_name == "Exponential":
                            distribution = Exponential(torch.tensor(case[0]))
                            k_tensor = torch.tensor(case[1], dtype=torch.float32)
                        elif dist_name == "LogNormal":
                            distribution = LogNormal(torch.tensor(case[0]), torch.tensor(case[1]))
                            k_tensor = torch.tensor(case[2], dtype=torch.float32)
                        elif dist_name == "Pareto":
                            distribution = Pareto(torch.tensor(case[0]), torch.tensor(case[1]))
                            k_tensor = torch.tensor(case[2], dtype=torch.float32)
                        elif dist_name == "Weibull":
                            distribution = Weibull(torch.tensor(1.0), torch.tensor(case[0]))
                            k_tensor = torch.tensor(case[1], dtype=torch.float32)
                        elif dist_name == "Dirichlet":
                            distribution = Dirichlet(case[0])
                            k_tensor = case[1]  # Dirichlet expects tensor input
                        elif dist_name == "VonMises":
                            distribution = VonMises(torch.tensor(case[0]), torch.tensor(case[1]))
                            k_tensor = torch.tensor(case[2], dtype=torch.float32)
                        elif dist_name == "StudentT":
                            distribution = StudentT(torch.tensor(case[0]))
                            k_tensor = torch.tensor(case[1], dtype=torch.float32)
                        elif dist_name == "Uniform":
                            distribution = Uniform(torch.tensor(case[0]), torch.tensor(case[1]))
                            k_tensor = torch.tensor(case[2], dtype=torch.float32)
                        else:
                            self.fail(f"Unhandled distribution: {dist_name}")

                        # Call wrapper log-prob method
                        if dist_name == "Dirichlet":
                            pdf = PtdmPdf().f(distribution, str(k_tensor.tolist()))[0]
                        else:
                            pdf = PtdmPdf().f(distribution, str(k_tensor.item()))[0]

                    # Ensure return values are tensors
                    self.assertTrue(isinstance(pdf, torch.Tensor), "Probability output is not a tensor")

                    # Convert tensors to floats for comparison
                    pdf = pdf.item()

                    # Check if computed values are approximately equal to expected values
                    self.assertAlmostEqual(pdf, case[-1], places=4, 
                                           msg=f"{dist_name}: Expected {case[-1]}, got {pdf}")

if __name__ == "__main__":
    unittest.main()