import sys
import os
import logging
import os.path as op

import numpy as num

from pyrocko import pile, trace, util, io
from pyrocko.parstack import parstack

from lassie import common, plot, grid as gridmod

logger = logging.getLogger('lassie.core')


def scan(
        config,
        override_tmin=None,
        override_tmax=None,
        show_detections=False,
        show_movie=False,
        force=False,
        stop_after_first=False):

    if config.detections_path:
        if op.exists(config.detections_path):
            if force:
                os.unlink(config.detections_path)
            else:
                raise common.LassieError(
                    'detections file already exists: %s' % config.detections_path)

        util.ensuredirs(config.detections_path)

    grid = config.get_grid()
    receivers = config.get_receivers()

    norm_map = gridmod.geometrical_normalization(grid, receivers)

    ifcs = config.image_function_contributions

    shift_tables = []
    tshift_maxs = []
    for ifc in ifcs:
        shift_tables.append(ifc.shifter.get_table(grid, receivers))
        tshift_maxs.append(shift_tables[-1].max())

    tshift_max = max(tshift_maxs) * 1.0

    tinc = tshift_max * 10

    tpad = max(ifc.get_tpad() for ifc in ifcs) + tshift_max

    fsmooth_min = min(ifc.get_fsmooth() for ifc in ifcs)

    blacklist = set(tuple(s.split('.')) for s in config.blacklist)

    station_index = dict(
        (rec.codes, i) for (i, rec) in enumerate(receivers)
        if rec.codes not in blacklist)

    for data_path in config.data_paths:
        if not op.exists(data_path):
            raise common.LassieError(
                'waveform data path does not exist: %s' % data_path)

    p = pile.make_pile(config.data_paths, fileformat='detect')
    if p.is_empty():
        raise common.LassieError('no usable waveforms found')

    ngridpoints = grid.size()

    idetection = 0

    station_weights = {}

    tmin = override_tmin or config.tmin
    tmax = override_tmax or config.tmax
    for trs in p.chopper(
            tmin=tmin, tmax=tmax, tinc=tinc, tpad=tpad,
            want_incomplete=False,
            trace_selector=lambda tr: tr.nslc_id[:3] in station_index):

        if not trs:
            continue

        logger.info('processing time window %s - %s' % (
            util.time_to_str(trs[0].wmin),
            util.time_to_str(trs[0].wmax)))

        wmin = trs[0].wmin
        wmax = trs[0].wmax

        frames = None
        pdata = []
        trs_debug = []
        shift_maxs = []
        for iifc, ifc in enumerate(ifcs):
            dataset = ifc.preprocess(trs, wmin, wmax, tshift_max)
            if not dataset:
                continue

            nstations_selected = len(dataset)

            nsls_selected, trs_selected = zip(*dataset)

            trs_debug.extend(trs + list(trs_selected))

            deltat_cf = trs_selected[0].deltat

            t0 = (wmin / deltat_cf) * deltat_cf

            istations_selected = num.array(
                [station_index[nsl] for nsl in nsls_selected], dtype=num.int)

            arrays = [tr.ydata.astype(num.float) for tr in trs_selected]

            offsets = num.array(
                [int(round((tr.tmin-t0) / deltat_cf)) for tr in trs_selected],
                dtype=num.int32)

            w = num.array(
                [station_weights.get(nsl, 1.0) for nsl in nsls_selected],
                dtype=num.float)


            weights = num.ones((ngridpoints, nstations_selected))
            weights *= w[num.newaxis, :]
            weights *= ifc.weight

            shift_table = shift_tables[iifc]

            shifts = -num.round(
                shift_table[:, istations_selected] /
                deltat_cf).astype(num.int32)

            pdata.append((list(trs_selected), shift_table, ifc))

            iwmin = int(round((wmin-t0) / deltat_cf))
            iwmax = int(round((wmax-t0) / deltat_cf))

            lengthout = iwmax - iwmin + 1

            shift_maxs.append(num.max(-shifts) * deltat_cf)

            frames, ioff = parstack(
                arrays, offsets, shifts, weights, 0,
                offsetout=iwmin,
                lengthout=lengthout,
                result=frames,
                impl='openmp')

        shift_max = max(shift_maxs)

        if config.sharpness_normalization:
            frame_maxs = frames.max(axis=0)
            frame_means = num.abs(frames).mean(axis=0)
            frames *= (frame_maxs / frame_means)[num.newaxis, :]
            frames *= norm_map[:, num.newaxis]

        frame_maxs = frames.max(axis=0)

        tmin_frames = t0 + ioff * deltat_cf

        tr_stackmax = trace.Trace(
            '', 'SMAX', '', '',
            tmin=tmin_frames,
            deltat=deltat_cf,
            ydata=frame_maxs)

        trs_debug.append(tr_stackmax)

        # trace.snuffle(trs_debug)

        ydata_window = tr_stackmax.chop(wmin, wmax, inplace=False).get_ydata()

        logger.info('CF stats: min %g, max %g, median %g' % (
            num.min(ydata_window),
            num.max(ydata_window),
            num.median(ydata_window)))

        tpeaks, apeaks = tr_stackmax.peaks(
            config.detector_threshold, shift_max + 1.0/fsmooth_min)

        for (tpeak, apeak) in zip(tpeaks, apeaks):
            if not (wmin <= tpeak and tpeak < wmax):
                continue

            logger.info('detection: %s %g' % (
                util.time_to_str(tpeak),
                apeak))

            iframe = int(round(((tpeak-t0) - ioff*deltat_cf) / deltat_cf))
            frame = frames[:, iframe]
            imax = num.argmax(frame)

            latpeak, lonpeak, xpeak, ypeak, zpeak = grid.index_to_location(imax)

            idetection += 1

            if config.detections_path:
                f = open(config.detections_path, 'a')
                f.write('%06i %s %g %g %g %g %g %g\n' % (
                    idetection,
                    util.time_to_str(tpeak, format='%Y-%m-%d %H:%M:%S.6FRAC'),
                    apeak,
                    latpeak, lonpeak, xpeak, ypeak, zpeak))

                f.close()

            if show_detections:
                fmin = min(ifc.fmin for ifc in ifcs)
                fmax = min(ifc.fmax for ifc in ifcs)
                plot.plot_detection(
                    grid, receivers, frames, tmin_frames,
                    deltat_cf, imax, iframe, fsmooth_min, xpeak, ypeak, zpeak,
                    tr_stackmax, tpeaks, apeaks, config.detector_threshold,
                    pdata, trs, fmin, fmax, idetection,
                    grid_station_shift_max=shift_max,
                    movie=show_movie)

            if stop_after_first:
                return
        
        tr_stackmax.chop(wmin, wmax)
        io.save([tr_stackmax], 'stackmax/trace_%(tmin_ms)s.mseed')


__all__ = [
    'scan',
]