import filecmp
import os
import unittest
import numpy as np

from simplekml import Link
from simplekml.featgeom import Feature, Geometry
from simplekml.timeprimitive import TimePrimitive

from eqep.data.eqdata import EQData
from eqep.interpolation.rbf_pgvinterpolator import RbfPGVInterpolator
from eqep.persisting.animation_saver import ShakeMapAnimationBuilder
from eqep.shakemap.display.scales import FullyTransparentSMScale
from eqep.shakemap.display.styles import PlainStyle
from eqep.shakemap.shakemap import ShakeMap


class AnimationTest(unittest.TestCase):
    """
    Unit-tests for `ShakeMapAnimationBuilder`
    """

    # directory with reference files
    REF_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'res',
                           'animation')

    # animation test parameters
    longs = np.array([15.523, 15.862, 17.671, 13.344, 16.395, 14.266, 17.543,
                      13.641, 14.549, 15.033])
    lats = np.array([47.250, 47.928, 48.859, 47.078, 49.062, 47.849, 49.777,
                     46.630, 46.509, 46.678])
    names = ['ARSA', 'CONA', 'JAVC', 'KBA', 'KRUC', 'MOA', 'MORC', 'MYKA',
             'OBKA', 'SOKA']
    values = np.array([33923, 76942, 18659, 5694, 89586, 15813, 58516, 2933,
                       7372, 7487])
    input = EQData(values, longs, lats, names)

    ALGS = ['multiquadric', 'inverse', 'gaussian', 'linear', 'cubic',
            'thin_plate']

    # parameters required for the tests
    dates = ['2015-02-15T12:00:01', '2015-02-15T12:00:02',
             '2015-02-15T12:00:03', '2015-02-15T12:00:04',
             '2015-02-15T12:00:05', '2015-02-15T12:00:06']
    images = ['cubic.png', 'gaussian.png', 'inverse.png', 'linear.png',
              'multiquadric.png', 'thin_plate.png']
    shakemaps = []

    @classmethod
    def setUpClass(cls):
        for i in range(len(cls.ALGS)):
            sm = ShakeMap(cls.input,
                          RbfPGVInterpolator(cls.ALGS[i]),
                          FullyTransparentSMScale(),
                          PlainStyle())

            sm.calc((13, 18), (46, 50))
            sm.calc_graphics((10, 8))

            sm.save_img(cls.images[i])
            cls.shakemaps.append(sm)

    @classmethod
    def tearDownClass(cls):
        for file in cls.images:
            os.remove(file)

    def setUp(self):
        #  reset the initial id-counters for deterministic tests
        Feature._id = 0
        TimePrimitive._id = 0
        Link._id = 0
        Geometry._id = 0

    def test_save_animation_kml(self):
        file = 'animation.kml'

        animation = ShakeMapAnimationBuilder()
        for sm, image, date in zip(self.shakemaps, self.images, self.dates):
            animation.add_shakemap(sm, image, date)

        animation.save(file)

        cmp = filecmp.cmp(file, os.path.join(self.REF_DIR, 'ref_animation.kml'))
        os.remove(file)

        self.assertTrue(cmp)

    def test_save_animation_kmz(self):
        file = 'animation.kmz'

        animation = ShakeMapAnimationBuilder()
        for sm, image, date in zip(self.shakemaps, self.images, self.dates):
            animation.add_shakemap(sm, image, date)

        animation.save(file)

        size = os.stat(file).st_size
        os.remove(file)

        self.assertEqual(os.stat(os.path.join(self.REF_DIR,
                                              'ref_animation.kmz')).st_size,
                         size)

    def test_save_animation_kml_random_order(self):
        file = 'animation.kml'

        random_indizes = [0, 5, 2, 4, 1, 3]

        animation = ShakeMapAnimationBuilder()
        for i in random_indizes:
            animation.add_shakemap(self.shakemaps[i], self.images[i],
                                   self.dates[i])

        animation.save(file)

        cmp = filecmp.cmp(file, os.path.join(self.REF_DIR,
                                             'ref_animation_random_order.kml'))
        os.remove(file)

        self.assertTrue(cmp)

    def test_save_animation_add_save(self):
        file = 'animation.kml'

        animation = ShakeMapAnimationBuilder()
        for sm, image, date in zip(self.shakemaps[:5], self.images, self.dates):
            animation.add_shakemap(sm, image, date)

        animation.save(file)

        # add another ShakeMap after saving and save() again
        animation.add_shakemap(self.shakemaps[5], self.images[5], self.dates[5])
        animation.save(file)

        cmp = filecmp.cmp(file, os.path.join(self.REF_DIR, 'ref_animation.kml'))
        os.remove(file)

        self.assertTrue(cmp)

    def test_save_animation_wrong_format(self):
        animation = ShakeMapAnimationBuilder()
        for sm, image, date in zip(self.shakemaps, self.images, self.dates):
            animation.add_shakemap(sm, image, date)

        self.assertRaises(ValueError, animation.save, 'animation.txt')


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