# http://pyrocko.org - GPLv3
#
# The Pyrocko Developers, 21st Century
# ---|P------/S----------~Lg----------
from __future__ import absolute_import, division, print_function

import os
import tarfile
import errno
import numpy as num
import time


from pyrocko.guts import Object, Timestamp
from pyrocko import gf, guts, util, pile, gmtpy

from .scenario import draw_scenario_gmt
from .error import ScenarioError

op = os.path
guts_prefix = 'pf.scenario'


def mtime(p):
    return os.stat(p).st_mtime


class ScenarioCollectionItem(Object):
    scenario_id = gf.StringID.T()
    time_created = Timestamp.T()

    def __init__(self, **kwargs):
        Object.__init__(self, **kwargs)
        self._path = None
        self._pile = None
        self._engine = None
        self._scenes = None

    def set_base_path(self, path):
        self._path = path

    def get_base_path(self):
        if self._path is None:
            raise EnvironmentError('Base path not set!')
        return self._path

    def init_modelling(self, engine):
        self._engine = engine

    def get_path(self, *entry):
        return op.join(*((self._path,) + entry))

    def get_generator(self):
        generator = guts.load(filename=self.get_path('generator.yaml'))
        generator.init_modelling(self._engine)
        return generator

    def get_time_range(self):
        return self.get_generator().get_time_range()

    def have_waveforms(self, tmin, tmax):
        p = self.get_waveform_pile()
        trs_have = p.all(
            tmin=tmin, tmax=tmax, load_data=False, degap=False)

        return any(tr.data_len() > 0 for tr in trs_have)

    def get_waveform_pile(self):
        self.ensure_data()

        if self._pile is None:
            path_waveforms = self.get_path('waveforms')
            util.ensuredir(path_waveforms)
            fns = util.select_files(
                [path_waveforms], show_progress=False)

            self._pile = pile.Pile()
            if fns:
                self._pile.load_files(
                    fns, fileformat='mseed', show_progress=False)

        return self._pile

    def get_insar_scenes(self):
        from kite import Scene
        if self._scenes is None:
            self._scenes = []
            path_insar = self.get_path('insar')
            util.ensuredir(path_insar)

            fns = util.select_files([path_insar], regex='\\.(npz)$',
                                    show_progress=False)
            for f in fns:
                self._scenes.append(Scene.load(f))

        return self._scenes

    def get_gnss_campaigns(self):
        return self.get_generator().get_gnss_campaigns()

    def make_map(self, path_pdf):
        draw_scenario_gmt(self.get_generator(), path_pdf)

    def get_map(self, format='pdf'):
        path_pdf = self.get_path('map.pdf')

        if not op.exists(path_pdf):
            self.make_map(path_pdf)

        path = self.get_path('map.%s' % format)

        outdated = op.exists(path) and mtime(path) < mtime(path_pdf)
        if not op.exists(path) or outdated:
            gmtpy.convert_graph(path_pdf, path)

        return path

    def ensure_data(self, tmin=None, tmax=None):
        return self.get_generator().ensure_data(
            self.get_path(), tmin, tmax)

    def get_archive(self):
        self.ensure_data()

        path_tar = self.get_path('archive.tar')
        if not op.exists(path_tar):
            path_base = self.get_path()
            path_waveforms = self.get_path('waveforms')
            self.ensure_data()

            fns = util.select_files(
                [path_waveforms], show_progress=False)

            f = tarfile.TarFile(path_tar, 'w')
            for fn in fns:
                fna = fn[len(path_base)+1:]
                f.add(fn, fna)

            f.close()

        return path_tar


class ScenarioCollection(object):

    def __init__(self, path, engine):
        self._scenario_suffix = 'scenario'
        self._path = path
        util.ensuredir(self._path)
        self._engine = engine
        self._load_scenarios()

    def _load_scenarios(self):
        scenarios = []
        base_path = self.get_path()
        for path_entry in os.listdir(base_path):
            scenario_id, suffix = op.splitext(path_entry)
            if suffix == '.' + self._scenario_suffix:
                path = op.join(base_path, path_entry, 'scenario.yaml')
                scenario = guts.load(filename=path)
                assert scenario.scenario_id == scenario_id
                scenario.set_base_path(op.join(base_path, path_entry))
                scenario.init_modelling(self._engine)
                scenarios.append(scenario)

        self._scenarios = scenarios
        self._scenarios.sort(key=lambda s: s.time_created)

    def get_path(self, scenario_id=None, *entry):
        if scenario_id is not None:
            return op.join(self._path, '%s.%s' % (
                scenario_id, self._scenario_suffix), *entry)
        else:
            return self._path

    def add_scenario(self, scenario_id, scenario_generator):

        if scenario_generator.seed is None:
            scenario_generator = guts.clone(scenario_generator)
            scenario_generator.seed = num.random.randint(1, 2**32-1)

        path = self.get_path(scenario_id)
        try:
            os.mkdir(path)
        except OSError as e:
            if e.errno == errno.EEXIST:
                raise ScenarioError(
                    'Scenario id is already in use: %s' % scenario_id)
            else:
                raise

        scenario = ScenarioCollectionItem(
            scenario_id=scenario_id,
            time_created=time.time())

        scenario_path = self.get_path(scenario_id, 'scenario.yaml')
        guts.dump(scenario, filename=scenario_path)

        generator_path = self.get_path(scenario_id, 'generator.yaml')
        guts.dump(scenario_generator, filename=generator_path)

        scenario.set_base_path(self.get_path(scenario_id))
        scenario.init_modelling(self._engine)

        self._scenarios.append(scenario)

    def list_scenarios(self, ilo=None, ihi=None):
        return self._scenarios[ilo:ihi]

    def get_scenario(self, scenario_id):
        for scenario in self._scenarios:
            if scenario.scenario_id == scenario_id:
                return scenario

        raise KeyError(scenario_id)
