# -*- coding: utf-8 -*-
"""
Created on Tue Jul 07 09:44:36 2015

@author: Marcus
"""

import pygimli as pg
import numpy as np
from pygimli.physics.traveltime.fastMarchingTest import fastMarch
import matplotlib.pyplot as plt
import time as time
from pygimli.mplviewer import drawMesh  # , drawField, drawStreamLines


class TravelTimeFMM(pg.ModellingBase):
    """
    Class that implements the Fast Marching Method (FMM). It can be used
    instead of Dijkstra modelling. Although it is currently quite slow!

    """
    def __init__(self, mesh, data, verbose=False):
        """
        Init function.

        Parameters:
        -----------
        mesh : pygimli.Mesh
            2D mesh to be used in the forward calculations.
        data : pygimli.DataContainer
            The datacontainer with sensor positions etc.
        verbose : boolean


            More printouts or not...
        """

        pg.ModellingBase.__init__(self, mesh, data, verbose)
        self.timefields = dict()
        self._jac = dict()
        self.num_sensors = data.sensorCount()
#        num_shots = len(np.unique(data("s")))

    def response(self, slowness):
        """
        Response function. Returns the result of the forwa
rd calculation.
        Uses the shot- and sensor positions specified in the data container.
        """

        mesh = self.mesh()
        param_markers = np.unique(mesh.cellMarker())
        param_count = len(param_markers)
        if len(slowness) == mesh.cellCount():
            self.mapModel(slowness)
        elif len(slowness) == param_count:
            # map the regions in the mesh to slowness
            slow_map = pg.stdMapF_F()
            min_reg_num = min(param_markers)
            for i, si in enumerate(slowness):
                slow_map.insert(float(i+min_reg_num), si)

            mesh.mapCellAttributes(slow_map)
        else:
            raise ValueError("Wrong no of parameters. Mesh size: {}, no "
                             "of regions: {}, and number of slowness values:"
                             "{}".format(self.mesh().cellCount(), param_count,
                                         len(slowness)))

        data = self.data()
        n_data = data.size()
        t_fmm = np.zeros(n_data)
        idx = 0
        for source_idx in [0]:  # np.unique(data("s")):
            # initialize source position and trvel time vector
            n_sensors = np.sum(data("s") == source_idx)
            # maybe not always same number of sensors
            source = data.sensorPosition(int(source_idx))
            times = pg.RVector(mesh.nodeCount(), 0.)

            # initialize sets and tags
#            upwind, downwind = set(), set()
            downwind = set()
            upTags = np.zeros(mesh.nodeCount())
            downTags = np.zeros(mesh.nodeCount())

            # define initial condition
            cell = mesh.findCell(source)

            for i, n in enumerate(cell.nodes()):
                times[n.id()] = cell.attribute() * n.pos().distance(source)
                upTags[n.id()] = 1
            for i, n in enumerate(cell.nodes()):
                tmpNodes = pg.commonNodes(n.cellSet())
                for nn in tmpNodes:
                    if not upTags[nn.id()] and not downTags[nn.id()]:
                        downwind.add(nn)
                        downTags[nn.id()] = 1

            # start fast marching
            while len(downwind) > 0:
                fastMarch(mesh, downwind, times, upTags, downTags)
            self.timefields[source_idx] = np.array(times)

            sensor_idx = data("g")[data("s") == source_idx]

            t_fmm[idx:idx+n_sensors] = np.array(
                [times[mesh.findNearestNode(data.sensorPosition(int(i)))]
                 for i in sensor_idx])
            idx += n_sensors

        return t_fmm

    def createJacobian(self, model):
        """
        Computes the jacobian matrix from the model.
        """
        pass

    def _intersect_lines(self, l1, l2):
        """
        Finds the parameters for which the two lines intersect.

        Assumes 2D lines!

        Parameters:
        -----------
        l1, l2 : pygimli Line
            Line objects from pygimli. Useful because they nicely wrap
            a line and has some utility functions.

        Returns:
        --------
        v : numpy array (length 2)
            The parameters (s and t) for l1 and l2, respectively. None if
            no intersection (i.e. parallell lines).
        """
#        print("l1: {}".format(l1))
#        print("l2: {}".format(l2))
#        first just check if parallell
        epsilon = 1.0e-4
        dir1 = l1.p1()-l1.p0()
        dir2 = l2.p1()-l2.p0()
#        print("dir1: {}, and length: {}".format(dir1, dir1.length()))
#        print("dir2: {}, and length: {}".format(dir2, dir2.length()))
        dir1 /= dir1.length()
        dir2 /= dir2.length()
#        print("dir1: {}, and length: {}".format(dir1, dir1.length()))
#        print("dir2: {}, and length: {}".format(dir2, dir2.length()))
        if abs(np.dot(dir1, dir2)) > 1.0 - epsilon:
            return np.array([1./epsilon, 1./epsilon])
#            raise Warning("parallell lines!")

        # Solve system: Av = w, were v, w are vectors. v = (s,t)
        ndim = 2
        A = np.ndarray((ndim, ndim))
        A[0, 0] = l1.p1().x() - l1.p0().x()
        A[1, 0] = l1.p1().y() - l1.p0().y()
        A[0, 1] = -(l2.p1().x() - l2.p0().x())
        A[1, 1] = -(l2.p1().y() - l2.p0().y())

        w = np.array([l2.p0().x() - l1.p0().x(), l2.p0().y() - l1.p0().y()])

        v = np.linalg.solve(A, w)
        if not np.allclose(np.dot(A, v), w):
            raise Warning("Problem with linear solver for intersection!")
        return v

    def _intersect_lines_by_points(self, p1, p2, p3, p4):
        """
        Finds the parameters for which the two lines intersect. The lines
        are defined by four points.

        Assumes 2D lines!

        Parameters:
        -----------
        p1, p2, p3, p4 : pygimli RVector3
            Position objects from pygimli. The lines are defined as:
                l1 : P1 to P2
                l2 : P3 to P4

        Returns:
        --------
        v : numpy array (length 2)
            The parameters (s and t) for l1 and l2, respectively. Will
            return "large" values if the lines are parallell.
        """

        # first just check if parallell
        epsilon = 1.0e-4
        dir1 = (p2 - p1).norm()
        dir2 = (p4 - p3).norm()

        if abs(np.dot(dir1, dir2)) > 1.0 - epsilon:
            return np.array([1./epsilon, 1./epsilon])
#            raise Warning("parallell lines!")

        # Solve system: Av = w, were v, w are vectors. v = (s,t)
        ndim = 2
        A = np.ndarray((ndim, ndim))
        A[0, 0] = p2.x() - p1.x()
        A[1, 0] = p2.y() - p1.y()
        A[0, 1] = p3.x() - p4.x()
        A[1, 1] = p3.y() - p4.y()

        w = np.array([p3.x() - p1.x(), p3.y() - p1.y()])

        v = np.linalg.solve(A, w)
        if not np.allclose(np.dot(A, v), w):
            raise Warning("Problem with linear solver for intersection!")
        return v[0], v[1]

    def _check_param(self, param, t_low=0, t_high=1.0):
        """
        Returns the "proper" t-value from the list. It should be
        positive along with the corresponding index value.
        """
        t_list = param[:, 0]
        par_pos = np.maximum(t_list, t_low)
        par_gt_eps = par_pos[par_pos > t_low+1e-5]
        print("t_list: {}\npar_pos: {}\npar_gt_eps: {}".format(
            t_list, par_pos, par_gt_eps))

        stay_on_edge = False
        try:
            t = np.min(par_gt_eps)
        except ValueError:
            stay_on_edge = True
            t = max(t_list)

        idx = int(param[t_list == t, 1][0])
        return t, idx, stay_on_edge

    def _check_param2(self, param, t_low=0, t_high=1.0):
        """
        Returns the "proper" t-v