#!/usr/bin/env python
# -*- coding:utf-8 -*-
from __future__ import unicode_literals

from datetime import datetime, timedelta
from multiprocessing.pool import Pool

from functools import partial

import numpy as np
from functools import wraps
from path import Path
try:
    from shutil import SameFileError
except ImportError:
    from shutil import Error as SameFileError
from tempfile import mkdtemp

from .tools import LagrantoException, run_cmd
from .formats import from_netcdf, to_ascii, from_ascii, to_netcdf

_all__ = ['Tra', 'LagrantoRun']


class Tra(object):
    """Class to work with LAGRANTO output.

    Read trajectories from a LAGRANTO file and return a structured numpy array

    Parameters
    ----------
        filename : string
            File containing lagranto trajectories
        usedatetime: bool
            Read times as datetime objects, default True
        array: structured array
            If defined creates a new Tra object filled with the array

    Returns
    -------
        structured array (Tra): trajs(ntra,ntime) with variables as tuple.

    Examples
    --------

    >>>    filename = 'mylslfile.nc'
    >>>    trajs = Tra()
    >>>    trajs.load_netcdf(filename)
    >>>    trajs['lon'][0,:]  # return the longitudes for the first trajectory.

    >>> trajs = Tra(filename)
    >>> selected_trajs = Tra(array=trajs[[10, 20, 30], :])

    Author : Nicolas Piaget, ETH Zurich , 2014
             Sebastiaan Crezee, ETH Zurich , 2014

    """

    _startdate = None

    def __init__(self, filename='', usedatetime=True, array=None, **kwargs):
        """Initialized a Tra object.

        If filename is given, try to load it directly;
        Arguments to the load function can be passed as key=value argument.
        """
        typefile = kwargs.pop('typefile', None)
        if typefile is not None:
            msg = 'typefile is not used anymore;' \
                  'it will be remove in futur version'
            raise DeprecationWarning(msg)
        if not filename:
            if array is None:
                self._array = None
            else:
                self._array = array
            return

        try:
            self.load_netcdf(filename, usedatetime=usedatetime, **kwargs)
        except (OSError, IOError, RuntimeError):
            try:
                self.load_ascii(filename, usedatetime=usedatetime, **kwargs)
            except:
                raise IOError("Unkown fileformat. Known formats \
                               are ascii or netcdf")

    def __len__(self):
        return len(self._array)

    def __getattr__(self, attr):
        if attr in self.__dict__:
            return getattr(self, attr)
        return getattr(self._array, attr)

    def __getitem__(self, key):
        return self._array[key]

    def __setitem__(self, key, item):
        if type(key) is slice:
            self._array = item
        elif key in self.dtype.names:
            self._array[key] = item
        else:
            dtypes = self._array.dtype.descr
            dtypes.append((key, item.dtype.descr[0][1]))
            dtypes = [(str(d[0]), d[1]) for d in dtypes]
            newarr = np.zeros(self._array.shape, dtype=dtypes)
            for var in self.variables:
                newarr[var] = self._array[var]
            newarr[key] = item
            self._array = newarr

    def __repr__(self):
        try:
            string = " \
            {} trajectories with {} time steps. \n \
            Available fields: {}\n \
            total duration: {} minutes".format(
                self.ntra, self.ntime,
                "/".join(self.variables),
                self.duration
            )
        except AttributeError:
            # Assume it's an empty Tra()
            string = "\
            Empty trajectories container.\n\
            Hint: use load_ascii() or load_netcdf()\n\
            to load data"
        return string

    @property
    def ntra(self):
        if self.ndim < 2:
            print(" \
                Be careful with the dimensions, \
                you may want to change the shape: \n \
                either shape + (1,) or (1,)+shape \
                ")
            return None
        return self.shape[0]

    @property
    def ntime(self):
        if self.ndim < 2:
            print(" \
                Be careful with the dimensions,\
                you may want to change the shape: \n \
                either shape + (1,) or (1,)+shape \
                ")
            return None
        return self.shape[1]

    @property
    def variables(self):
        return list(self.dtype.names)

    @property
    def duration(self):
        """Time duration in minutes."""
        end = self['time'][0, -1]
        start = self['time'][0, 0]
        delta = end - start
        if type(delta) == np.timedelta64:
            return delta.astype(timedelta).total_seconds() / 60.
        else:
            return delta * 60.

    @property
    def initial(self):
        """Give the initial time of the trajectories."""
        starttime = self['time'][0, 0]
        return starttime.astype(datetime)

    @property
    def startdate(self):
        if self._startdate is None:
            time0 = self['time'][0, 0]
            if type(time0) == np.datetime64:
                self._startdate = time0.astype(datetime)
            else:
                self._startdate = datetime(1900, 1, 1, 0)
        return self._startdate

    def set_array(self, array):
        """To change the trajectories array."""
        self._array = array

    def concatenate(self, trajs, time=False, inplace=False):
        """To concatenate trajectories together.

        Concatenate trajectories together and return a new object.
        The trajectories should contain the same variables.
        if time=False, the number of timestep in each trajs should be the same
        if time=True, the number of trajectories in each trajs should be the same

        Parameters
        ----------

            trajs: Tra or list of Tra
                Trajectories to concatenate with the current one
            time: bool, default False
                if True concatenate along the time dimension
            inplace: bool, default False
                if True append the trajs to current Tra object and return None

        Returns
        -------
        Tra
            Return a new Tra (trajectories) object

        Examples
        --------

        Create a new Tra object which include trajs and all newtrajs

        >>> trajs = Tra(filename)
        >>> newtrajs = [Tra(f) for f in files]
        >>> alltrajs = trajs.concatenate(newtrajs)

        Append newtrajs to trajs

        >>> trajs = Tra(filename)
        >>> newtrajs = [Tra(f) for f in files]
        >>> trajs.concatenate(newtrajs, inplace=True)

        """
        if type(trajs) not in [tuple, list]:
            trajs = (trajs,)
        if time:
            trajstuple = (self._array.T,)
            trajstuple += tuple(tra._array.T for tra in trajs)
            test = np.concatenate(trajstuple).T
        else:
            trajstuple = (self._array,)
            trajstuple += tuple(tra._array for tra in trajs)
            test = np.concatenate(trajstuple)

        if inplace:
            self._array = test
        else:
            newtrajs = Tra()
            newtrajs.set_array(test)
            return newtrajs

    def append(self, trajs):
        """append trajectories

        Parameters
        ---------
        trajs: single Tra or list of Tra
                Trajectories to concatenate with the current one

        Examples
        --------

        >>> trajs = Tra(filename)
        >>> newtrajs = [Tra(f) for f in files]
        >>> trajs.append(newtrajs)
        """
        self.concatenate(trajs, inplace=True)

    def write(self, filename, fileformat='netcdf'):
        globals()['_write_{}'.format(fileformat)](self, filename)

    def load_netcdf(self, filename, usedatetime=True, msv=-999, unit='hours',
                    exclude=None, date=None):
        self._array, self._startdate = from_netcdf(filename,
                                                   usedatetime=usedatetime,
                                                   msv=msv,
                                                   unit=unit,
                                                   exclude=exclude,
                                                   date=date)
    load_netcdf.__doc__ = from_netcdf.__doc__

    def write_netcdf(self, filename, exclude=None, unit='hours'):
        to_netcdf(self, filename, exclude=exclude, unit=unit)
    write_netcdf.__doc__ = to_netcdf.__doc__

    def write_ascii(self, filename, gz=False):
        to_ascii(self, filename, gz=gz)
    write_ascii.__doc__ = to_ascii.__doc__

    def load_ascii(self, filename, usedatetime=True, msv=-999.999, gz=False):
        self._array, self._startdate = from_ascii(filename,
                                                  usedatetime=usedatetime,
                                                  msv=msv,
                                                  gz=gz)
    load_ascii.__doc__ = from_ascii.__doc__


def intmpdir(f):
    @wraps(f)
    def wrapper(*args, **kwds):
        workingdir = args[0]
        no_tmp_dir = kwds.get('no_tmp_dir', False)
        if Path(workingdir).parent == '/tmp' or no_tmp_dir:
            return f(*args, **kwds)
        else:
            tmpdir = Path(mkdtemp())
            lwkds = {key: value for key, value in kwds.items()
                     if key in ['outputdir', 'version', 'sdate']}
            link_files(workingdir, tmpdir, **lwkds)
            nargs = (tmpdir,) + args[1:]
            out = f(*nargs, **kwds)
            tmpdir.rmdir_p()
            return out
    return wrapper


@intmpdir
def create_startf(workingdir, date, filename, specifier, version='cosmo',
                  outputdir='', sdate=None, tolist=False, cmd_header=None,
                  options=None):
    """Create a startfile for LAGRANTO."""
    outputdir = Path(outputdir)
    workingdir = Path(workingdir)
    create = "startf.{version} "
    create += "{date:%Y%m%d_%H} {filename} {specifier}"
    if options is not None:
        if type(options) is list:
            for option in options:
                create += ' ' + option
        else:
            raise LagrantoException('options as to be a list of string')
    # if tolist:
    #     filename = Path(filename).splitext()[0] + '.4'
    create = create.format(version=version, date=date,
                           filename=filename, specifier=specifier)
    out = run_cmd(create, workingdir, cmd_header=cmd_header)
    if tolist:
        lslname = filename
        filename = Path(lslname).splitext()[0] + '.startf'
        lsl2list = 'lsl2list.{} {} {}'.format(version, lslname, filename)
        out += run_cmd(lsl2list, workingdir, cmd_header=cmd_header)
    try:
        Path.copy(workingdir / filename, outputdir / filename)
    except SameFileError:
        pass
    if tolist:
        return out, filename
    else:
        return out


@intmpdir
def select(workingdir, inpfile, outfile, crit,
           outputdir='', sdate=None, version='cosmo',
           cmd_header=None, **kwargs):
    """Select trajectories."""
    netcdf_format = kwargs.pop('netcdf_format', 'CF')
    select_cmd = 'select.{version} {inpfile} {outfile} "{crit}"'
    select_cmd = select_cmd.format(version=version, inpfile=inpfile,
                                   outfile=outfile, crit=crit)
    out = run_cmd(select_cmd, workingdir,
                  netcdf_format=netcdf_format, cmd_header=cmd_header)
    outputdir = Path(outputdir)
    try:
        Path.copy(workingdir / outfile, outputdir / outfile)
    except SameFileError:
        pass
    return out


@intmpdir
def trace(workingdir, filename, outputdir='', outfile='', tracevars='',
          tracevars_content='', field='', sdate=None, version='cosmo',
          cmd_header=None, **kwargs):
    """Trace variables along a trajectory."""
    netcdf_format = kwargs.pop('netcdf_format', 'CF')
    workingdir = Path(workingdir)
    outputdir = Path(outputdir)
    if not outfile:
        outfile = filename
    trace_cmd = 'trace.{version} {filename} {outfile}'
    tracevars_file = 'tracevars'
    if tracevars:
        tracevars_file = tracevars
        trace_cmd += ' -v ' + tracevars
    if tracevars_content:
        with (Path(workingdir) / tracevars_file).open('w') as f:
            f.write(tracevars_content)
    if field:
        trace_cmd += ' -f ' + field
    trace_cmd = trace_cmd.format(filename=filename, outfile=outfile,
                                 version=version)
    out = run_cmd(trace_cmd, workingdir,
                  netcdf_format=netcdf_format, cmd_header=cmd_header)
    try:
        Path.copy(workingdir / outfile, outputdir / outfile)
    except SameFileError:
        pass
    return out


@intmpdir
def caltra(workingdir, startdate, enddate, startfile, filename,
           jump=True, outputdir='', sdate=None, version='cosmo',
           cmd_header=None, withtrace=False, tracevars_content='',
           tracevars='', **kwargs):
    """Calculate trajectories for air parcels.

    starting at positions specified in startfile.
    """
    outputdir = Path(outputdir)
    netcdf_format = kwargs.pop('netcdf_format', 'CF')
    no_tmp_dir = kwargs.pop('no_tmp_dir', False)
    prefix = kwargs.pop('prefix', '')
    if no_tmp_dir:
        filename = outputdir / filename
    if not withtrace:    
        caltra_cmd = 'caltra.{version}'
        caltra_cmd += ' {startdate:%Y%m%d_%H} {enddate:%Y%m%d_%H}'
        caltra_cmd += ' {startfile} {filename}'
        if jump:
            caltra_cmd += ' -j'
        for key, value in kwargs.items():
            caltra_cmd += ' -{key} {value}'.format(key=key, value=value)
        caltra_cmd = caltra_cmd.format(startdate=startdate, enddate=enddate,
                                       startfile=startfile, filename=filename,
                                       version=version)
    else:
        tracevars_file = tracevars if tracevars else 'tracevars'
        if tracevars_content:
            with (Path(workingdir) / tracevars_file).open('w') as f:
                f.write(tracevars_content)
        delta = (enddate - startdate).total_seconds() / (60 * 60)
        caltra_cmd = 'caltra {startfile} {delta:1.0f} {filename} '
        caltra_cmd += '-ref {startdate:%Y%m%d_%H%M}'
        for key, value in kwargs.items():
            caltra_cmd += ' -{key} {value}'.format(key=key, value=value)
        caltra_cmd = caltra_cmd.format(startdate=startdate, delta=delta,
                                       startfile=startfile, filename=filename)
    out = run_cmd(caltra_cmd, workingdir, netcdf_format=netcdf_format,
                  cmd_header=cmd_header, prefix=prefix)
    if not no_tmp_dir:
        try:
            Path.copy(workingdir / filename, outputdir / filename)
        except SameFileError:
            pass
    return out


@intmpdir
def density(workingdir, inpfile, outfile,
            outputdir='', sdate=None, version='cosmo',
            cmd_header=None, **kwargs):
    """Calculate trajectories density."""
    outputdir = Path(outputdir)
    netcdf_format = kwargs.pop('netcdf_format', 'CF')
    density_cmd = 'density.{version} {inpfile} {outfile}'
    for key, value in kwargs.items():
        density_cmd += ' -{key} {value}'.format(key=key, value=value)
    density_cmd = density_cmd.format(inpfile=inpfile, outfile=outfile,
                                     version=version)
    out = run_cmd(density_cmd, workingdir,
                  netcdf_format=netcdf_format, cmd_header=cmd_header)
    try:
        Path.copy(workingdir / outfile, outputdir / outfile)
    except SameFileError:
        pass
    return out


class LagrantoRun:
    """Perform Lagranto calculation.

    Parameters
    ----------
    dates: list
        list of (startdate, enddate) tuple
    workingdir: string, optional
        path to the model output directory, default to current
    outputdir: string, optional
        path to the trajectory utput directory, defautl to current
    startf: string, optional
        name of the startf to use (or to create), default to startf.4
    lslname: string, optional
        name of the lsl file, define its type, default to lsl_{:%Y%m%d%H}.4
    tracevars: string, optional
        name of a tracevars file as used by trace, default to none
    field: string, optional
        name of a single field to trace, default to none
    version: string, optional
        name of the model version to use, currently only cosmo (default)
    linkfiles: function, optional
        function used to overwrite link_files in run.
        Should be used if COSMO output is not standard netcdf
    nprocesses: int, optional
        Number of processes used when running in parallel, default to 10
    sdate: datetime object,
        Starting date of the simulation;
        useful if files are named in forecast mode
    fresh: bool, optional
        Fresh start. Remove output directory first.
    cmd_header: string, optional
        Change the header using for running the command;
        see ``run_cmd`` help for more details.
    """

    def __init__(self, dates, workingdir='.', outputdir='.',
                 startf='startf.4', lslname='lsl_{:%Y%m%d%H}.4',
                 tracevars='', field='', version='cosmo', linkfiles=None,
                 nprocesses=10, sdate=None, fresh=False, cmd_header=None):
        self.dates = dates
        self.workingdir = Path(workingdir).abspath()
        self.outputdir = Path(outputdir).abspath()
        if fresh:
            self.clear()
        self.outputdir.makedirs_p()
        self.startf = Path(startf)
        self.lslname = lslname
        self.tracevars = tracevars
        self.field = field
        self.version = version
        self.link_files = linkfiles if linkfiles else link_files
        self.nprocesses = nprocesses
        self.sdate = sdate
        self.cmd_header = cmd_header

    def create_startf(self, date, specifier, filename='', tolist=False,
                      **kwargs):
        """Create a file with the starting position of the trajectories.

        Parameters
        ----------
        date: datetime
            date for which the startf date is produced
        specifier: string
            detailed description of starting positions;
            (see LAGRANTO documentations for more details)
        filename: string
            filename where starting positions are saved;
            the default is defined by `LagrantoRun`
        tolist: bool
            If the starting position should be saved as a list of points;
             Useful if the same file is used for different starting times
        kwargs: dict
            key, value options, possible values are listed bellow;

            - cmd_header:
              header to add to the shell command;
              (see `run_cmd` for more details)
            - options:
              list of options to pass to the startf function;
              (see the LAGRANTO documentation for more details)


        """
        if filename:
            self.startf = Path(filename)
        out = create_startf(self.workingdir, date, self.startf,
                            specifier, outputdir=self.outputdir,
                            sdate=self.sdate, version=self.version,
                            tolist=tolist, cmd_header=self.cmd_header, **kwargs)
        if tolist:
            self.startf = out[1]
            return out[0]
        else:
            return out

    def caltra(self, startdate, enddate, filename='', **kwargs):
        if filename:
            self.lslname = filename
        self.caltra_out = self.lslname.format(startdate)
        return caltra(self.workingdir, startdate, enddate, self.startf,
                      self.caltra_out, outputdir=self.outputdir,
                      sdate=self.sdate, version=self.version,
                      cmd_header=self.cmd_header, **kwargs)

    def trace(self, date, filename='', tracevars='', tracevars_content='',
              field='', **kwargs):
        r"""Trace variable along a trajectory.

        Trace meteorological fields along trajectories

        Parameters
        ----------
        date: datetime object
            starting date of the trajectories to trace
        filename: string, optional
            If not specified use `LagrantoRun.lslname`
        tracevars: string, optional
            Name of the file with the field to trace;
            If not specified use `LagrantoRun.tracevars`
        tracevars_content: string, optional
            Content of the tracevars file;

            .. code_block: python

            \"\"\"\n
            QV      1000. P 1\n
            PV      1.    T 1\n
            \"\"\"

        field: string,
            Specify a field to trace as follow: QV 1.
        kwargs
        """
        tracef = tracevars if tracevars else self.tracevars
        fieldv = field if field else self.field
        filename = filename if filename else self.lslname
        filename = filename.format(date)
        return trace(self.workingdir, filename, outputdir=self.outputdir,
                     tracevars=tracef, tracevars_content=tracevars_content,
                     field=fieldv, sdate=self.sdate,
                     version=self.version, cmd_header=self.cmd_header, **kwargs)

    def density(self, inpfile=None, outfile=None, **kwargs):
        if inpfile is None:
            try:
                inpfile = self.caltra_out
            except AttributeError:
                raise ValueError('Provide an input file or run caltra')
        if outfile is None:
            self.density_out = 'density.4'
        else:
            self.density_out = outfile
        return density(self.workingdir, inpfile=inpfile,
                       outfile=self.density_out, outputdir=self.outputdir,
                       sdate=self.sdate, version=self.version,
                       cmd_header=self.cmd_header, **kwargs)

    def select(self, inpfile=None, outfile=None, **kwargs):
        if inpfile is None:
            try:
                inpfile = self.caltra_out
            except AttributeError:
                raise ValueError('Provide an input file or run caltra')
        if outfile is None:
            self.select_out = 'selected.4'
        else:
            self.select_out = outfile
        return select(self.workingdir, inpfile=inpfile,
                      outfile=self.select_out, outputdir=self.outputdir,
                      sdate=self.sdate, version=self.version,
                      cmd_header=self.cmd_header, **kwargs)

    def run(self, caltra_kw=None, trace_kw=None, **kwargs):
        """Run caltra and trace.

        if kwargs are provided they are passed to the link_files function
        """
        if caltra_kw is None:
            caltra_kw = dict()
        if trace_kw is None:
            trace_kw = dict()
        for sd, ed in self.dates:
            self.single_run(sd, ed, caltra_kw=caltra_kw, trace_kw=trace_kw,
                            **kwargs)

    def run_parallel(self, caltra_kw=None, trace_kw=None, **kwargs):
        """Run caltra and trace in parallel.

        Similar to run() but using multiprocessing.Pool
        """
        if caltra_kw is None:
            caltra_kw = dict()
        if trace_kw is None:
            trace_kw = dict()

        single_run = partial(self.single_run, caltra_kw=caltra_kw,
                             trace_kw=trace_kw, **kwargs)
        with Pool(processes=min(self.nprocesses, len(self.dates))) as pool:
            results = pool.starmap(single_run, self.dates)
        return results

    def single_run(self, sd, ed, caltra_kw=None, trace_kw=None, type='both',
                   debug=False, **kwargs):
        if caltra_kw is None:
            caltra_kw = dict()
        if trace_kw is None:
            trace_kw = dict()

        workingdir = self.workingdir
        tmpdir = Path(mkdtemp())
        self.workingdir = tmpdir
        nkwargs = {'version': self.version,
                   'outputdir': self.outputdir,
                   'sdate': self.sdate}
        nkwargs.update(kwargs)
        self.link_files(workingdir, tmpdir, **nkwargs)
        try:
            if type == 'both':
                out = self.caltra(sd, ed, **caltra_kw)
                out += self.trace(sd, **trace_kw)
            elif type == 'caltra':
                out = self.caltra(sd, ed, **caltra_kw)
            elif type == 'trace':
                out = self.trace(sd, **trace_kw)
        except LagrantoException as err:
            err.args += (tmpdir,)
            if debug:
                raise
            else:
                return err
        self.workingdir = workingdir
        tmpdir.rmtree_p()
        return out

    def clear(self):
        self.outputdir.rmtree_p()


def link_files(folder, tmpdir, outputdir='', version='cosmo',
               sdate=None, sfiles=True):
    """Link data from <folder> to <tmpdir> as P,S files.

    Parameters
    ----------
    folder: string,
        path to the origin folder
    tmpdir: string,
        path to the destination folder
    outputdir: string,
        path to the folder where results are saved
    version: string,
        only cosmo is available for now
    sdate: datetime object,
        useful is the cosmo file are named in forecast mode
    """
    folder = Path(folder).abspath()
    tmpdir = Path(tmpdir)
    files = set(folder.files())
    if version == 'cosmo':
        constant_file = folder.files('l*c.nc')
        if constant_file:
            constant_file[0].symlink(tmpdir / 'LMCONSTANTS')
        cosmofiles = set(folder.files('lff[df]*[0-9].nc'))
        for fname in cosmofiles:
                try:
                    date = datetime.strptime(fname.name, 'lffd%Y%m%d%H.nc')
                except ValueError as err:
                    sfname = fname.name[4:-3]
                    dd, hh, mm, ss = [int(sfname[i:i + 2])
                                      for i in range(0, len(sfname), 2)]
                    date = sdate + timedelta(days=dd, hours=hh,
                                             minutes=hh, seconds=ss)
                pfile = 'P{:%Y%m%d_%H}'.format(date)
                fname.symlink(tmpdir / pfile)
                if sfiles:
                    sfile = 'S{:%Y%m%d_%H}'.format(date)
                    fname.symlink(tmpdir / sfile)
        files = files.difference(cosmofiles)
    for fname in files:
        try:
            fname.symlink(tmpdir / fname.name)
        except FileExistsError:
            pass
    if outputdir:
        outputdir = Path(outputdir).abspath()
        for fname in outputdir.files():
            try:
                fname.symlink(tmpdir / fname.name)
            except FileExistsError:
                pass
