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

import os
import math
import logging

import numpy as num

from pyrocko import trace, util, plot
from pyrocko.guts import Object, Int, String, Timestamp

from . import io_common

logger = logging.getLogger('pyrocko.io.datacube')

N_GPS_TAGS_WANTED = 200  # must match definition in datacube_ext.c


def color(c):
    c = plot.color(c)
    return tuple(x/255. for x in c)


class DataCubeError(io_common.FileLoadError):
    pass


class ControlPointError(Exception):
    pass


def make_control_point(ipos_block, t_block, tref, deltat):

    # reduce time (no drift would mean straight line)
    tred = (t_block - tref) - ipos_block*deltat

    # first round, remove outliers
    q25, q75 = num.percentile(tred, (25., 75.))
    iok = num.logical_and(q25 <= tred, tred <= q75)

    # detrend
    slope, offset = num.polyfit(ipos_block[iok], tred[iok], 1)
    tred2 = tred - (offset + slope * ipos_block)

    # second round, remove outliers based on detrended tred, refit
    q25, q75 = num.percentile(tred2, (25., 75.))
    iok = num.logical_and(q25 <= tred2, tred2 <= q75)
    x = ipos_block[iok].copy()
    ipos0 = x[0]
    x -= ipos0
    y = tred[iok].copy()
    (slope, offset), cov = num.polyfit(x, y, 1, cov=True)

    slope_err, offset_err = num.sqrt(num.diag(cov))
    slope_err_limit = 1.0e-10
    offset_err_limit = 5.0e-3

    if slope_err > slope_err_limit:
        raise ControlPointError('slope error too large')

    if offset_err > offset_err_limit:
        raise ControlPointError('offset error too large')

    ic = ipos_block[ipos_block.size//2]
    tc = offset + slope * (ic - ipos0)

    return ic, tc + ic * deltat + tref


def analyse_gps_tags(header, gps_tags, offset, nsamples):

    ipos, t, fix, nsvs = gps_tags
    deltat = 1.0 / int(header['S_RATE'])

    tquartz = offset + ipos * deltat

    toff = t - tquartz
    toff_median = num.median(toff)

    n = t.size

    dtdt = (t[1:n] - t[0:n-1]) / (tquartz[1:n] - tquartz[0:n-1])

    ok = abs(toff_median - toff) < 10.

    xok = num.abs(dtdt - 1.0) < 0.00001

    ok[0] = False
    ok[1:n] &= xok
    ok[0:n-1] &= xok
    ok[n-1] = False

    ipos = ipos[ok]
    t = t[ok]
    fix = fix[ok]
    nsvs = nsvs[ok]

    blocksize = N_GPS_TAGS_WANTED // 2

    try:
        if ipos.size < blocksize:
            raise ControlPointError(
                'could not safely determine time corrections from gps')

        j = 0
        control_points = []
        tref = num.median(t - ipos*deltat)
        while j < ipos.size - blocksize:
            ipos_block = ipos[j:j+blocksize]
            t_block = t[j:j+blocksize]
            try:
                ic, tc = make_control_point(ipos_block, t_block, tref, deltat)
                control_points.append((ic, tc))
            except ControlPointError:
                pass
            j += blocksize

        ipos_last = ipos[-blocksize:]
        t_last = t[-blocksize:]
        try:
            ic, tc = make_control_point(ipos_last, t_last, tref, deltat)
            control_points.append((ic, tc))
        except ControlPointError:
            pass

        if len(control_points) < 2:
            raise ControlPointError(
                'could not safely determine time corrections from gps')

        i0, t0 = control_points[0]
        i1, t1 = control_points[1]
        i2, t2 = control_points[-2]
        i3, t3 = control_points[-1]
        if len(control_points) == 2:
            tmin = t0 - i0 * deltat - offset * deltat
            tmax = t3 + (nsamples - i3 - 1) * deltat
        else:
            icontrol = num.array(
                [x[0] for x in control_points], dtype=num.int64)
            tcontrol = num.array(
                [x[1] for x in control_points], dtype=num.float)
            # robust against steps:
            slope = num.median(
                (tcontrol[1:] - tcontrol[:-1])
                / (icontrol[1:] - icontrol[:-1]))

            tmin = t0 + (offset - i0) * slope
            tmax = t2 + (offset + nsamples - 1 - i2) * slope

        if offset < i0:
            control_points[0:0] = [(offset, tmin)]

        if offset + nsamples - 1 > i3:
            control_points.append((offset + nsamples - 1, tmax))

        icontrol = num.array([x[0] for x in control_points], dtype=num.int64)
        tcontrol = num.array([x[1] for x in control_points], dtype=num.float)

        return tmin, tmax, icontrol, tcontrol, ok

    except ControlPointError:

        tmin = util.str_to_time(header['S_DATE'] + header['S_TIME'],
                                format='%y/%m/%d%H:%M:%S')

        idat = int(header['DAT_NO'])
        if idat == 0:
            tmin = tmin + util.gps_utc_offset(tmin)
        else:
            tmin = util.day_start(tmin + idat * 24.*3600.) \
                + util.gps_utc_offset(tmin)

        tmax = tmin + (nsamples - 1) * deltat
        icontrol, tcontrol = None, None
        return tmin, tmax, icontrol, tcontrol, None


def plot_timeline(fns):
    from matplotlib import pyplot as plt

    fig = plt.figure()
    axes = fig.gca()

    h = 3600.

    if isinstance(fns, str):
        fn = fns
        if os.path.isdir(fn):
            fns = [
                os.path.join(fn, entry) for entry in sorted(os.listdir(fn))]

            ipos, t, fix, nsvs, header, offset, nsamples = \
                get_timing_context(fns)

        else:
            ipos, t, fix, nsvs, header, offset, nsamples = \
                get_extended_timing_context(fn)

    else:
        ipos, t, fix, nsvs, header, offset, nsamples = \
            get_timing_context(fns)

    deltat = 1.0 / int(header['S_RATE'])

    tref = num.median(t - ipos * deltat)
    tref = round(tref / deltat) * deltat

    x = ipos*deltat
    y = (t - tref) - ipos*deltat

    bfix = fix != 0
    bnofix = fix == 0

    tmin, tmax, icontrol, tcontrol, ok = analyse_gps_tags(
        header, (ipos, t, fix, nsvs), offset, nsamples)

    la = num.logical_and
    nok = num.logical_not(ok)

    axes.plot(
        x[la(bfix, ok)]/h, y[la(bfix, ok)], '+',
        ms=5, color=color('chameleon3'))
    axes.plot(
        x[la(bfix, nok)]/h, y[la(bfix, nok)], '+',
        ms=5, color=color('aluminium4'))

    axes.plot(
        x[la(bnofix, ok)]/h, y[la(bnofix, ok)], 'x',
        ms=5, color=color('chocolate3'))
    axes.plot(
        x[la(bnofix, nok)]/h, y[la(bnofix, nok)], 'x',
        ms=5, color=color('aluminium4'))

    tred = tcontrol - icontrol*deltat - tref
    axes.plot(icontrol*deltat/h, tred, color=color('aluminium6'))
    axes.plot(icontrol*deltat/h, tred, 'o', ms=5, color=color('aluminium6'))

    ymin = (math.floor(tred.min() / deltat)-1) * deltat
    ymax = (math.ceil(tred.max() / deltat)+1) * deltat

    # axes.set_ylim(ymin, ymax)
    if ymax - ymin < 1000 * deltat:
        ygrid = math.floor(tred.min() / deltat) * deltat
        while ygrid < ymax:
            axes.axhline(ygrid, color=color('aluminium4'), alpha=0.3)
            ygrid += deltat

    xmin = icontrol[0]*deltat/h
    xmax = icontrol[-1]*deltat/h
    xsize = xmax - xmin
    xmin -= xsize * 0.1
    xmax += xsize * 0.1
    axes.set_xlim(xmin, xmax)

    axes.set_ylim(ymin, ymax)

    axes.set_xlabel('Uncorrected (quartz) time [h]')
    axes.set_ylabel('Relative time correction [s]')

    plt.show()


g_dir_contexts = {}


class DirContextEntry(Object):
    path = String.T()
    tstart = Timestamp.T()
    ifile = Int.T()


class DirContext(Object):
    path = String.T()
    mtime = Timestamp.T()
    entries = DirContextEntry.T()

    def get_entry(self, fn):
        path = os.path.abspath(fn)
        for entry in self.entries:
            if entry.path == path:
                return entry

        raise Exception('entry not found')

    def iter_entries(self, fn, step=1):
        current = self.get_entry(fn)
        by_ifile = dict(
            (entry.ifile, entry) for entry in self.entries
            if entry.tstart == current.tstart)

        icurrent = current.ifile
        while True:
            icurrent += step
            try:
                yield by_ifile[icurrent]

            except KeyError:
                break


def context(fn):
    from pyrocko import datacube_ext

    dpath = os.path.dirname(os.path.abspath(fn))
    mtimes = [os.stat(dpath)[8]]

    dentries = sorted([os.path.join(dpath, f) for f in os.listdir(dpath)
                       if os.path.isfile(os.path.join(dpath, f))])
    for dentry in dentries:
        fn2 = os.path.join(dpath, dentry)
        mtimes.append(os.stat(fn2)[8])

    mtime = float(max(mtimes))

    if dpath in g_dir_contexts:
        dir_context = g_dir_contexts[dpath]
        if dir_context.mtime == mtime:
            return dir_context

        del g_dir_contexts[dpath]

    entries = []
    for dentry in dentries:
        fn2 = os.path.join(dpath, dentry)
        if not os.path.isfile(fn2):
            continue

        with open(fn2, 'rb') as f:
            first512 = f.read(512)
            if not detect(first512):
                continue

        with open(fn2, 'rb') as f:
            try:
                header, data_arrays, gps_tags, nsamples, _ = \
                        datacube_ext.load(f.fileno(), 3, 0, -1, None)

            except datacube_ext.DataCubeError as e:
                e = DataCubeError(str(e))
                e.set_context('filename', fn)
                raise e

        header = dict(header)
        entries.append(DirContextEntry(
            path=os.path.abspath(fn2),
            tstart=util.str_to_time(
                '20' + header['S_DATE'] + ' ' + header['S_TIME'],
                format='%Y/%m/%d %H:%M:%S'),
            ifile=int(header['DAT_NO'])))

    dir_context = DirContext(mtime=mtime, path=dpath, entries=entries)

    return dir_context


def get_time_infos(fn):
    from pyrocko import datacube_ext

    with open(fn, 'rb') as f:
        try:
            header, _, gps_tags, nsamples, _ = datacube_ext.load(
                f.fileno(), 1, 0, -1, None)

        except datacube_ext.DataCubeError as e:
            e = DataCubeError(str(e))
            e.set_context('filename', fn)
            raise e

    return dict(header), gps_tags, nsamples


def get_timing_context(fns):
    joined = [[], [], [], []]
    ioff = 0
    for fn in fns:
        header, gps_tags, nsamples = get_time_infos(fn)

        ipos = gps_tags[0]
        ipos += ioff

        for i in range(4):
            joined[i].append(gps_tags[i])

        ioff += nsamples

    ipos, t, fix, nsvs = [num.concatenate(x) for x in joined]

    nsamples = ioff
    return ipos, t, fix, nsvs, header, 0, nsamples


def get_extended_timing_context(fn):
    c = context(fn)

    header, gps_tags, nsamples_base = get_time_infos(fn)

    ioff = 0
    aggregated = [gps_tags]

    nsamples_total = nsamples_base

    if num.sum(gps_tags[2]) == 0:

        ioff = nsamples_base
        for entry in c.iter_entries(fn, 1):

            _, gps_tags, nsamples = get_time_infos(entry.path)

            ipos = gps_tags[0]
            ipos += ioff

            aggregated.append(gps_tags)
            nsamples_total += nsamples

            if num.sum(gps_tags[2]) > 0:
                break

            ioff += nsamples

        ioff = 0
        for entry in c.iter_entries(fn, -1):

            _, gps_tags, nsamples = get_time_infos(entry.path)

            ioff -= nsamples

            ipos = gps_tags[0]
            ipos += ioff

            aggregated[0:0] = [gps_tags]

            nsamples_total += nsamples

            if num.sum(gps_tags[2]) > 0:
                break

    ipos, t, fix, nsvs = [num.concatenate(x) for x in zip(*aggregated)]

#    return ipos, t, fix, nsvs, header, ioff, nsamples_total
    return ipos, t, fix, nsvs, header, 0, nsamples_base


def iload(fn, load_data=True, interpolation='sinc'):
    from pyrocko import datacube_ext
    from pyrocko import signal_ext

    if interpolation not in ('sinc', 'off'):
        raise NotImplementedError(
            'no such interpolation method: %s' % interpolation)

    with open(fn, 'rb') as f:
        if load_data:
            loadflag = 2
        else:
            loadflag = 1

        try:
            header, data_arrays, gps_tags, nsamples, _ = datacube_ext.load(
                f.fileno(), loadflag, 0, -1, None)

        except datacube_ext.DataCubeError as e:
            e = DataCubeError(str(e))
            e.set_context('filename', fn)
            raise e

    header = dict(header)
    deltat = 1.0 / int(header['S_RATE'])
    nchannels = int(header['CH_NUM'])

    ipos, t, fix, nsvs, header_, offset_, nsamples_ = \
        get_extended_timing_context(fn)

    tmin, tmax, icontrol, tcontrol, _ = analyse_gps_tags(
        header_, (ipos, t, fix, nsvs), offset_, nsamples_)

    if icontrol is None:
        logger.warn(
            'No usable GPS timestamps found. Using datacube header '
            'information to guess time. (file: "%s")' % fn)

    tmin_ip = round(tmin / deltat) * deltat
    if interpolation != 'off':
        tmax_ip = round(tmax / deltat) * deltat
    else:
        tmax_ip = tmin_ip + (nsamples-1) * deltat

    nsamples_ip = int(round((tmax_ip - tmin_ip)/deltat)) + 1
    # to prevent problems with rounding errors:
    tmax_ip = tmin_ip + (nsamples_ip-1) * deltat

    leaps = num.array(
        [x[0] + util.gps_utc_offset(x[0]) for x in util.read_leap_seconds2()],
        dtype=num.float)

    if load_data and icontrol is not None:
        ncontrol_this = num.sum(
            num.logical_and(0 <= icontrol, icontrol < nsamples))

        if ncontrol_this <= 1:
            logger.warn(
                'Extrapolating GPS time information from directory context '
                '(insufficient number of GPS timestamps in file: "%s").' % fn)

    for i in range(nchannels):
        if load_data:
            arr = data_arrays[i]
            assert arr.size == nsamples

            if interpolation == 'sinc' and icontrol is not None:

                ydata = num.empty(nsamples_ip, dtype=num.float)
                try:
                    signal_ext.antidrift(
                        icontrol, tcontrol,
                        arr.astype(num.float), tmin_ip, deltat, ydata)

                except signal_ext.Error as e:
                    e = DataCubeError(str(e))
                    e.set_context('filename', fn)
                    e.set_context('n_control_points', icontrol.size)
                    e.set_context('n_samples_raw', arr.size)
                    e.set_context('n_samples_ip', ydata.size)
                    e.set_context('tmin_ip', util.time_to_str(tmin_ip))
                    raise e

                ydata = num.round(ydata).astype(arr.dtype)
            else:
                ydata = arr

            tr_tmin = tmin_ip
            tr_tmax = None
        else:
            ydata = None
            tr_tmin = tmin_ip
            tr_tmax = tmax_ip

        tr = trace.Trace('', header['DEV_NO'], '', 'p%i' % i, deltat=deltat,
                         ydata=ydata, tmin=tr_tmin, tmax=tr_tmax, meta=header)

        bleaps = num.logical_and(tmin_ip <= leaps, leaps < tmax_ip)

        if num.any(bleaps):
            assert num.sum(bleaps) == 1
            tcut = leaps[bleaps][0]

            for tmin_cut, tmax_cut in [
                    (tr.tmin, tcut), (tcut, tr.tmax+tr.deltat)]:

                try:
                    tr_cut = tr.chop(tmin_cut, tmax_cut, inplace=False)
                    tr_cut.shift(
                        util.utc_gps_offset(0.5*(tr_cut.tmin+tr_cut.tmax)))
                    yield tr_cut

                except trace.NoData:
                    pass

        else:
            tr.shift(util.utc_gps_offset(0.5*(tr.tmin+tr.tmax)))
            yield tr


header_keys = {
    str: b'GIPP_V DEV_NO E_NAME GPS_PO S_TIME S_DATE DAT_NO'.split(),
    int: b'''P_AMPL CH_NUM S_RATE D_FILT C_MODE A_CHOP F_TIME GPS_TI GPS_OF
            A_FILT A_PHAS GPS_ON ACQ_ON V_TCXO D_VOLT E_VOLT'''.split()}

all_header_keys = header_keys[str] + header_keys[int]


def detect(first512):
    s = first512

    if len(s) < 512:
        return False

    if ord(s[0:1]) >> 4 != 15:
        return False

    n = s.find(b'\x80')
    if n == -1:
        n = len(s)

    s = s[1:n]
    s = s.replace(b'\xf0', b'')
    s = s.replace(b';', b' ')
    s = s.replace(b'=', b' ')
    kvs = s.split(b' ')

    if len([k for k in all_header_keys if k in kvs]) == 0:
        return False
    return True


if __name__ == '__main__':
    import sys
    fns = sys.argv[1:]
    if len(fns) > 1:
        plot_timeline(fns)
    else:
        plot_timeline(fns[0])
