import os, re, time
import flopy
import sys
import numpy as np
import pandas as pd
import pyvista as pv
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import matplotlib.patheffects as PathEffects
from scipy.interpolate import griddata
from mf6Voronoi.utils import isRunningInJupyter, printBannerHtml, printBannerText

def FlowVectorGenerator(gwf, backgroundImageDict=None, 
                        kstpkper=(0,0),
                        plotGrid = True,
                        plotContour = True,
                        contourLevels = 10,
                        layer = 0,
                        istep = 4,
                        jstep = 4,
                        scale = 10,
                        normalize=True
                        ):
    """
    Mf6Voronoi utility for the representation of flow direction vectors in a Matplotlib figure. The function works not only with normal Modflow6 Dis models like the generated by Model Muse but also with Voronoi based Modflow6 Disv models. There is a capability to insert background images and other functionalities, however the styling options are closed in order to preserve the artistic style of the plot.

    Parameters
    ----------
    bakgroundImageDisct : dict
        Dictionary with the shape of {'fig':path,''wrl':path} pointing to the paths of the image and world file.
    kstpkper :  tuple of ints
        A tuple containing the time step and stress period. Only the last time step of the stress period is plotted. 
    plotGrid : bool
        Option to represent heads grid for a given layer on a grid.
    plotContour : bool
        Option to represent head contours for a given layer. 
    contourLevels : int
        Number of contours for the plotContour representation.
    layer : integer
        MODFLOW zero-based layer number for the head and contour representation.
    istep : integer
        Row frequency to plot (default is 4).
    jstep : integer
        Row frequency to plot (default is 4).
    scale: float
        Scales the length of the arrow inversely.    
    normalize : bool
        Boolean flag used to determine if vectors should be normalized.
    """

    if backgroundImageDict != None:
        try: 
            figPath = backgroundImageDict['fig'] 
            wrlPath = backgroundImageDict['wrl']
        except KeyError:
            print('ERROR: Something went wrong with the background image')
        #open figure and world file:
        bcgImg = mpimg.imread(figPath)
        wrlFile = open(wrlPath).read().split('\n')
        cellsize = float(wrlFile[0])
        left = float(wrlFile[4])
        right = float(wrlFile[4]) + bcgImg.shape[1]*cellsize
        bottom = float(wrlFile[5]) - bcgImg.shape[0]*cellsize
        top = float(wrlFile[5])      

    # Build geometry vtk
    if gwf.modelgrid.grid_type == 'structured' or gwf.modelgrid.grid_type == 'vertex':
        dis = gwf.get_package('dis')
        try:

            head = gwf.output.head().get_data(kstpkper)
            bud = gwf.output.budget()
            spdis = bud.get_data(text='DATA-SPDIS')[kstpkper[1]]

            hactive = head[head>gwf.hdry]
            hactive = head[head<gwf.hnoflo]
            levels = np.linspace(hactive.min(),hactive.max(),contourLevels)

            fig, ax = plt.subplots(figsize=(10, 10)) ## Org

            qx, qy, qz = flopy.utils.postprocessing.get_specific_discharge(spdis, gwf)

            arrowColor = 'steelblue'

            #initialize graph
            pmv = flopy.plot.PlotMapView(gwf, layer=layer)
            #add layers
            pmv.plot_grid(colors='turquoise', lw=0.3, alpha=0.5,ax=ax, zorder=1)
            
            if plotGrid:
                pmv.plot_array(head, masked_values=[1e+30], cmap='YlGnBu',
                               ax=ax, zorder=2, alpha=0.5)
                arrowColor = 'azure'
            elif backgroundImageDict:
                arrowColor = 'azure'

            if plotContour:
                mf6Contour = pmv.contour_array(head, levels=levels, linewidths=3.,
                                               cmap='YlGnBu',
                                               ax=ax, zorder=3)
                texts = plt.clabel(mf6Contour, inline=True,
                                   fontsize=8, fmt="%.1f")

                for text in texts:
                    text.set_color('deepskyblue')
                    text.set_path_effects([
                        PathEffects.withStroke(linewidth=2, foreground='white')
                    ])
            pmv.plot_vector(qx, qy, normalize=normalize, color=arrowColor, 
                            scale_units='width',scale=scale, 
                            istep=istep,
                            jstep=jstep,
                            ax=ax, zorder=4)
            if backgroundImageDict != None:
                image = ax.imshow(bcgImg, extent=(left, right, bottom, top),
                                  zorder=0, alpha=0.7)
                
            return fig


        except AttributeError:
            print('ERROR: Something went wrong with model data, check the dis file')

        
    # elif gwf.modelgrid.grid_type == 'vertex':
    #     pass
    elif gwf.modelgrid.grid_type == 'unstructured':
        print("DISU: This dicretization type is not supported")
    else:
        print("No Dis file was found")

    