import unittest
from pymatgen.core import Structure
import pandas as pd
import numpy as np
from veidt.monte_carlo.ensemble import NVT, NPT, uVT
from veidt.monte_carlo.base import State, StaticState
from veidt.monte_carlo.state import AtomNumberState, IsingState
from veidt.monte_carlo.base import StateDict
from veidt.monte_carlo.state import SpinStructure
from veidt.abstract import Model, Describer
import os

file_path = os.path.dirname(__file__)


class SimpleLinearModel(Model):
    def __init__(self, describer):
        self.describer = describer

    def fit(self):
        pass

    def predict(self, objs):
        features = self.describer.transform(objs)
        return np.sum(5*features**2, axis=1)


class NaCount(Describer):
    def describe(self, structure):
        return pd.DataFrame({'Na': [np.sum([i.specie.name == 'Na' for i in structure])],
                             'K': [np.sum([i.specie.name == 'K' for i in structure])]})


class Volume(State):
    def __init__(self, state, name='volume'):
        super(Volume, self).__init__(state, name)
        self.label = 'volume'

    def change(self):
        self.state += (np.random.rand(1) - 0.5) * 0.1 * self.state


class TestEnsemble(unittest.TestCase):

    def setUp(self):
        self.structure = Structure.from_file(os.path.join(file_path, 'test_NaCoO2.cif'))

    def test_nvt(self):
        model = SimpleLinearModel(NaCount())
        nvt = NVT(model)
        state_dict = StateDict([StaticState(100, 'temperature'),
                                AtomNumberState(10),
                                IsingState([0]*22+[1, 1])])
        spin_struct = SpinStructure(structure=self.structure, species_map={1: "Na", 0: "K"},
                                    state_dict=state_dict)
        energy1 = nvt.exponential(spin_struct)
        self.assertEqual(energy1, 2440)
        spin_struct.change()
        energy2 = nvt.exponential(spin_struct)
        self.assertIn(energy2, (2250, 2650))

    def test_npt(self):
        model = SimpleLinearModel(NaCount())
        npt = NPT(model)
        state_dict = StateDict([StaticState(100, 'temperature'),
                                Volume(10, 'volume'),
                                StaticState(-1, 'pressure'),
                                AtomNumberState(10),
                                IsingState([0, 1]*12)])
        spin_struct = SpinStructure(structure=self.structure, species_map={1: "Na", 0: "K"},
                                    state_dict=state_dict)
        energy1 = npt.exponential(spin_struct)
        self.assertAlmostEqual(energy1, 1429.80157, 4)
        spin_struct.change()
        energy2 = npt.exponential(spin_struct)
        self.assertLessEqual(energy2, 1429.80157 + 10 + 5)
        self.assertLessEqual(1429.80157 + 10 - 5, energy2)

    def test_uVT(self):
        model = SimpleLinearModel(NaCount())
        uvt = uVT(model)
        state_dict = StateDict([StaticState(300, 'temperature'),
                                StaticState(10, 'volume'),
                                StaticState(23, 'm'),
                                StaticState(-3, 'mu'),
                                StaticState(-1, 'pressure'),
                                AtomNumberState(10),
                                IsingState([0, 1]*12)])
        spin_struct = SpinStructure(structure=self.structure, species_map={1: "Na", 0: "K"},
                                    state_dict=state_dict)
        energy1 = uvt.exponential(spin_struct)
        #self.assertAlmostEqual(energy1, 1469.6704985, 4)

        spin_struct2 = spin_struct.copy()
        spin_struct2.change()
        print(sum(spin_struct.state_dict['ising'].state), sum(spin_struct2.state_dict['ising'].state))
        print(uvt.d_exponential(spin_struct, spin_struct2))


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

