#!/usr/bin/env python


__author__ = "waroquiers"

import json
import os
import unittest

from pymatgen.analysis.chemenv.coordination_environments.chemenv_strategies import (
    AngleNbSetWeight,
    CNBiasNbSetWeight,
    DeltaCSMNbSetWeight,
    DeltaDistanceNbSetWeight,
    DistanceAngleAreaNbSetWeight,
    DistanceNbSetWeight,
    NormalizedAngleDistanceNbSetWeight,
    SelfCSMNbSetWeight,
)
from pymatgen.analysis.chemenv.coordination_environments.structure_environments import (
    StructureEnvironments,
)
from pymatgen.util.testing import PymatgenTest

se_files_dir = os.path.join(
    os.path.dirname(__file__),
    "..",
    "..",
    "..",
    "..",
    "..",
    "test_files",
    "chemenv",
    "structure_environments_files",
)


class FakeNbSet:
    def __init__(self, cn=None):
        self.cn = cn

    def __len__(self):
        return self.cn

    pass


class DummyStructureEnvironments:
    pass


class DummyVoronoiContainer:
    pass


class StrategyWeightsTest(PymatgenTest):
    def test_angle_weight(self):
        fake_nb_set = FakeNbSet()
        dummy_se = DummyStructureEnvironments()

        # Angles for a given fake nb_set with 5 neighbors
        fake_nb_set.angles = [
            1.8595833644514066,
            2.622518848090717,
            3.08570351705799,
            2.2695472184920042,
            2.2695338778592387,
        ]
        angle_weight = AngleNbSetWeight(aa=1.0)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 0.9634354419021528, delta=1e-8)
        angle_weight = AngleNbSetWeight(aa=2.0)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 0.92820785071319645, delta=1e-8)
        angle_weight = AngleNbSetWeight(aa=0.5)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 0.98154747307613843, delta=1e-8)

        self.assertNotEqual(AngleNbSetWeight(1.0), AngleNbSetWeight(2.0))

        # nb_set with no neighbor
        fake_nb_set.angles = []
        angle_weight = AngleNbSetWeight(aa=1.0)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 0.0, delta=1e-8)
        angle_weight = AngleNbSetWeight(aa=2.0)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 0.0, delta=1e-8)

        # nb_set with one neighbor
        fake_nb_set.angles = [3.08570351705799]
        angle_weight = AngleNbSetWeight(aa=1.0)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 0.24555248382791284, delta=1e-8)
        angle_weight = AngleNbSetWeight(aa=2.0)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 0.060296022314057396, delta=1e-8)
        angle_weight = AngleNbSetWeight(aa=0.5)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 0.49553252549950022, delta=1e-8)

        # nb_set with 6 neighbors (sum of the angles is 4*pi, i.e. the full sphere)
        fake_nb_set.angles = [
            1.8595833644514066,
            0.459483788407816,
            2.622518848090717,
            3.08570351705799,
            2.2695472184920042,
            2.2695338778592387,
        ]
        angle_weight = AngleNbSetWeight(aa=1.0)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 1.0, delta=1e-8)
        angle_weight = AngleNbSetWeight(aa=2.0)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 1.0, delta=1e-8)
        angle_weight = AngleNbSetWeight(aa=0.5)
        aw = angle_weight.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(aw, 1.0, delta=1e-8)

    def test_normalized_angle_distance_weight(self):
        fake_nb_set = FakeNbSet()
        dummy_se = DummyStructureEnvironments()

        nadw1 = NormalizedAngleDistanceNbSetWeight(average_type="geometric", aa=1, bb=1)
        nadw2 = NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=1, bb=1)
        nadw3 = NormalizedAngleDistanceNbSetWeight(average_type="geometric", aa=0, bb=1)
        nadw4 = NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=1, bb=0)
        nadw5 = NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=0.1, bb=0.1)
        nadw6 = NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=0, bb=0.1)
        nadw7 = NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=0.1, bb=0)
        nadw8 = NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=2, bb=0)
        nadw9 = NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=0, bb=2)
        nadw10 = NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=2, bb=2)
        nadw11 = NormalizedAngleDistanceNbSetWeight(average_type="geometric", aa=1, bb=2)
        nadw12 = NormalizedAngleDistanceNbSetWeight(average_type="geometric", aa=2, bb=1)
        self.assertNotEqual(nadw11, nadw12)
        with self.assertRaisesRegex(ValueError, "Both exponents are 0."):
            NormalizedAngleDistanceNbSetWeight(average_type="arithmetic", aa=0, bb=0)
        with self.assertRaisesRegex(
            ValueError,
            'Average type is "arithmetix" ' 'while it should be "geometric" or "arithmetic"',
        ):
            NormalizedAngleDistanceNbSetWeight(average_type="arithmetix", aa=1, bb=1)

        fake_nb_set.normalized_distances = [
            1.2632574171572457,
            1.1231971151388764,
            1.0,
            1.1887986376446249,
            1.188805134890625,
        ]
        fake_nb_set.normalized_angles = [
            0.6026448601336767,
            0.8498933334305273,
            1.0,
            0.7355039801931018,
            0.7354996568248028,
        ]
        w1 = nadw1.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w1, 0.67310887189488189, delta=1e-8)
        w2 = nadw2.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w2, 0.69422258996523023, delta=1e-8)
        w3 = nadw3.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w3, 0.8700949310182079, delta=1e-8)
        w4 = nadw4.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w4, 0.7847083661164217, delta=1e-8)
        w5 = nadw5.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w5, 0.96148050989126843, delta=1e-8)
        w6 = nadw6.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w6, 0.98621181678741754, delta=1e-8)
        w7 = nadw7.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w7, 0.97479580875402994, delta=1e-8)
        w8 = nadw8.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w8, 0.63348507114489783, delta=1e-8)
        w9 = nadw9.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w9, 0.7668954450583646, delta=1e-8)
        w10 = nadw10.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w10, 0.51313920014833292, delta=1e-8)
        w11 = nadw11.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w11, 0.585668617459, delta=1e-8)
        w12 = nadw12.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w12, 0.520719679281, delta=1e-8)

    def test_CN_bias_weight(self):
        fake_nb_set = FakeNbSet()
        dummy_se = DummyStructureEnvironments()
        bias_weight1 = CNBiasNbSetWeight.linearly_equidistant(weight_cn1=1.0, weight_cn13=13.0)
        bias_weight2 = CNBiasNbSetWeight.geometrically_equidistant(weight_cn1=1.0, weight_cn13=1.1 ** 12)
        bias_weight3 = CNBiasNbSetWeight.explicit(
            cn_weights={
                1: 1.0,
                2: 3.0,
                3: 3.2,
                4: 4.0,
                5: 4.1,
                6: 4.2,
                7: 4.3,
                8: 4.4,
                9: 4.5,
                10: 4.6,
                11: 4.6,
                12: 4.7,
                13: 4.8,
            }
        )
        with self.assertRaisesRegex(ValueError, "Weights should be provided for CN 1 to 13"):
            CNBiasNbSetWeight.explicit(cn_weights={1: 1.0, 13: 2.0})

        fake_nb_set.cn = 1
        w1 = bias_weight1.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w1, 1.0, delta=1e-8)
        w2 = bias_weight2.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w2, 1.0, delta=1e-8)
        w3 = bias_weight3.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w3, 1.0, delta=1e-8)
        fake_nb_set.cn = 7
        w1 = bias_weight1.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w1, 7.0, delta=1e-8)
        w2 = bias_weight2.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w2, 1.1 ** 6, delta=1e-8)
        w3 = bias_weight3.weight(nb_set=fake_nb_set, structure_environments=dummy_se)
        self.assertAlmostEqual(w3, 4.3, delta=1e-8)
        fake_nb_set.cn = 13