"""
Set up the plot figures, axes, and items to be done for each frame.
This module is imported by the plotting routines and then the
function setplot is called to set the plot parameters.

modified 12 December, 2025 by SeanPaul La Selle, slaselle@usgs.gov

This software is preliminary or provisional and is subject to revision. It is
being provided to meet the need for timely best science. The software has not
received final approval by the U.S. Geological Survey (USGS). No warranty,
expressed or implied, is made by the USGS or the U.S. Government as to the
functionality of the software and related material nor shall the fact of release
constitute any such warranty. The software is provided on the condition that
neither the USGS nor the U.S. Government shall be held liable for any damages
resulting from the authorized or unauthorized use of the software.
"""

import pylab
import glob, os, sys
from importlib import reload
from numpy import loadtxt, cos, pi,linspace,zeros,ones,hstack
from matplotlib import image
import matplotlib
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource
import numpy as np
import numpy
from clawpack.visclaw import gridtools
from clawpack.geoclaw import fgout_tools, topotools, dtopotools, marching_front
from pylab import plot,legend,xlabel,ylabel,sqrt,grid,xlim,ylim,\
                    nan,where,title,fill_between,xticks,yticks,text

from clawpack.pyclaw import Solution



outdir= os.path.abspath('_output')

# for transect through source and westport
ymin, ymax = 5259500, 5258715
xmin, xmax = 434140, 434740
xtrans = linspace(xmin, xmax, 1000)
ytrans = linspace(ymin, ymax, 1000)

#--------------------------
def setplot(plotdata=None):
#--------------------------

    """
    Specify what is to be plotted at each frame.
    Input:  plotdata, an instance of pyclaw.plotters.data.ClawPlotData.
    Output: a modified version of plotdata.

    """
    from clawpack.visclaw import colormaps, geoplot
    if plotdata is None:
        from clawpack.visclaw.data import ClawPlotData, ClawPlotItem

        plotdata_dry = ClawPlotData(outdir=outdir)
        plotdata_dry.outdir = outdir
        plotdata_dry.format="binary"
        plotdata_dry.clearfigures()

    plotdata.clearfigures()  # clear any old figures,axes,items data

    # To plot gauge locations on pcolor or contour plot, use this as
    # an afteraxis function:

    def addgauges(current_data):
        from pyclaw.plotters import gaugetools
        gaugetools.plot_gauge_locations(current_data.plotdata, \
             gaugenos='all', format_string='ko', add_labels=True)
    def topo_q(current_data):
        q=current_data.q
        b=current_data.aux
        var = q[0,:,:]+b[0,:,:]
        return var

    def slide_only(current_data):
        q=current_data.q
        var = numpy.ma.masked_where(q[0,:,:]>0.,q[0,:,:])
        return var

    def topo_wo_slide(current_data):
        h=current_data.q[0,:,:]
        b=current_data.aux[0,:,:]
        var = numpy.ma.masked_where(h>0.,b)
        return var

    def topo(current_data):
        b=current_data.aux
        var = b[0,:,:]
        return var

    def surface_or_depth(current_data):
        from numpy import ma, where
        q = current_data.q
        h = q[0,:,:]
        eta = q[6,:,:]
        topo = eta - h
        surface = ma.masked_where(h<=drytol, eta)
        depth = ma.masked_where(h<=drytol, h)
        surface_or_depth =  where(topo<0, surface, depth)
        return surface_or_depth

    def hillshade(array, azimuth, angle_altitude):

        from numpy import gradient
        from numpy import pi
        from numpy import arctan
        from numpy import arctan2
        from numpy import sin
        from numpy import cos
        from numpy import sqrt
        #from numpy import zeros
        #from numpy import uint8

        x, y = gradient(array)#,dx,dy)
        slope = pi/2. - arctan(sqrt(x*x + y*y))
        aspect = arctan2(-x, y)
        azimuthrad = azimuth*pi / 180.
        altituderad = angle_altitude*pi / 180.
        shaded = sin(altituderad) * sin(slope) + cos(altituderad) * cos(slope) * cos(azimuthrad - aspect)

        return 255*(shaded + 1)/2

    #-----------------------------------------
    #-----------------------------------------
    #-----------------------------------------
    # transition pcolor
    #-----------------------------------------
    #-----------------------------------------
    #-----------------------------------------
    plotfigure = plotdata.new_plotfigure(name='Domain', figno=0)
    plotfigure.kwargs = {'figsize': (8,5)}
    plotfigure.show = True

    # Set up for axes in this figure:
    plotaxes = plotfigure.new_plotaxes()
    plotaxes.axescmd = 'axes([0.05,0.66,.3,.33])' 
    plotaxes.title = 'McCormick Slide h:m:s'
    plotaxes.scaled = True
    plotaxes.xlimits = 'auto'
    plotaxes.ylimits = [5258600 , 5259600]
    plotaxes.xticks_kwargs = {'rotation':20}


    # Landslide
    plotitem = plotaxes.new_plotitem(plot_type='2d_pcolor')
    colors = plt.cm.get_cmap('RdPu', 10)
    plotitem.plot_var = 0
    plotitem.pcolor_cmap = colors
    plotitem.pcolor_cmin =  0.
    plotitem.pcolor_cmax =  5.
    plotitem.add_colorbar = True
    plotitem.amr_celledges_show = [0,0,0]
    plotitem.patchedges_show = 0
    plotitem.show = True


    def aa(current_data):
        pd = current_data.plotdata
        frameno = current_data.frameno
        ax = plt.gca()
        t = current_data.t
        t = t / 60.  # minutes
        x = current_data.x
        y = current_data.y
        h = current_data.q[0,:,:]
        b = current_data.aux[0,:,:]
        tmp = numpy.where(h==0.,b,numpy.nan)
        # bathy
        ax.contour(x,y,b, np.arange(-100, 100, 5),
                    colors='k', 
                    alpha=0.5, 
                    negative_linestyles='solid',
                    linewidths=0.2)
        
        # shoreline contour
        ax.contour(x,y,b, [57], colors='k', 
                   linewidths=2)
        # #topo
        # ax.contour(dem.X, dem.Y, dem.Z, np.arange(0, 3000, 20),
        #             colors='k',
        #             alpha=0.5)
        ax.ticklabel_format(useOffset=False)
        plt.xticks(fontsize=10,rotation=20)
        plt.yticks(fontsize=10)
       
        # transect line
        ax.plot([xmin, xmax], [ymin, ymax], 'ko-', alpha=0.5)
    
    plotaxes.afteraxes = aa

    # ________________
    # Transition Transect Axis
    # ________________
    plotaxes = plotfigure.new_plotaxes('transect')
    plotaxes.axescmd = 'axes([0.5,.66,.5,.33])'
    

    def plot_xsec(current_data):
        pd = current_data.plotdata
        frameno = current_data.frameno
        framesoln = Solution(frameno, path=outdir, file_format=pd.format)
       
        t = current_data.t
        t = t / 60.  # minutes
        x = current_data.x
        y = current_data.y
        h = current_data.q[0,:,:]
        b = current_data.aux[0,:,:]
        tmp = numpy.where(h==0.,b,numpy.nan)
        # eta_trans = gridtools.grid_output_2d(framesoln, -1, xtrans, ytrans)
        h_trans = gridtools.grid_output_2d(framesoln, 0, xtrans, ytrans)
        B_trans = gridtools.grid_output_2d(framesoln, -1, xtrans, ytrans)
        slide_base = B_trans-h_trans
        fill_between(xtrans, -100*ones(xtrans.shape), slide_base, color=[.8,1,.8])
        fill_between(xtrans, slide_base, B_trans, color='purple')
        plot(xtrans, B_trans, 'k')
        plot(xtrans, slide_base, 'k', alpha=0.5, label='slide surface')
        plt.hlines(57,xmin, xmax, colors='b', linestyles='-', label='water level')


        #xlabel('distance',fontsize=10)
        ylabel('surface elevation (m)',fontsize=10)
        xlabel('UTM Easting', fontsize=10)
        xlim(xmin,xmax)
        ylim(-10-5,70-5)
        grid(True)
        title('Along transect (%d, %d) to (%d, %d)' %(xmin, ymin, xmax, ymax),fontsize=12)
        legend(loc='upper right', fontsize=8)
        xticks(fontsize=10,rotation=90)
        yticks(fontsize=10)
        # x1 = -124.15
        # x2 = x1 + 2000/(111e3*cos(pi*47/180))
        # y1 = -7
        # plot([x1,x2],[y1,y1],'k',lw=3)
        # text(0.5*(x1+x2), y1+0.5, '2 km', fontsize=10, ha='center',va='bottom')
        

    plotaxes.afteraxes = plot_xsec

    
    # Parameters used only when creating html and/or latex hardcopy
    # e.g., via clawpack.visclaw.frametools.printframes:

    plotdata.printfigs = True                # print figures
    plotdata.print_format = 'png'            # file format
    #plotdata.print_framenos = np.arange(30,61,1).tolist()        # list of frames to print
    # plotdata.print_framenos = 'all'
    # plotdata.print_framenos = [1,5,10,15,20,25,30,45,50,55,60, 65, 70, 75]
    plotdata.print_fignos = 'all'            # list of figures to print
    plotdata.print_gaugenos = 'all'          # list of gauges to print
    plotdata.html = True                     # create html files of plots?
    plotdata.html_homelink = '../README.html'   # pointer for top of index
    plotdata.latex = True                    # create latex file of plots?
    plotdata.latex_figsperline = 2           # layout of plots
    plotdata.latex_framesperline = 1         # layout of plots
    plotdata.latex_makepdf = False           # also run pdflatex?
    plotdata.parallel = True

    return plotdata

    
