"""Loads and/or plots 2D, topologlically structured data on quadrilaterals using matplotlib.
"""

import sys,os
import numpy as np
import h5py
import mesh
import colors
import functools

def fullname(varname):
    fullname = varname
    if not '.cell.' in fullname:
        fullname = fullname+'.cell.0'
    return fullname


def transect_data(varnames, keys='all', directory=".", filename="ats_vis_data.h5",
                  mesh_filename="ats_vis_mesh.h5", coord_order=None, deformable=False, return_map=False):
    """Pulls simulation output into structured 2D arrays for transect-based, (i,j) indexing.

    Input:
      varnames       | A list of variable names to pull, e.g.
                     |  ['saturation_liquid', 'saturation_ice'], or a single variable
                     |  name, e.g. 'saturation_liquid'
      keys           | Indices of timesteps to pull.  Either an int (i.e. 0, -1, etc) 
                     |  for the kth timestep, or a list of ints, or 'all'.
      directory      | Directory of the run.  Defaults to '.'
      filename       | Filename of the run.  Defaults to 'visdump_data.h5'
      mesh_filename  | Filename of the mesh.  Defaults to 'visdump_mesh.h5'
      coord_order    | Order of the transect coordinates.  Defaults to ['x','z'].  The 
                     |  mesh is sorted in this order.
      deformable     | Is the mesh deforming?
      return_map     | See return value below.
   
    Output:
      Output is an array of shape:
      ( len(varnames+2), len(keys), n_cells_coord_order[0], n_cells_coord_order[1] )
    
      data[0,0,:,:] is the coord_order[0] centroid
      data[1,0,:,:] is the coord_order[1] centroid
      data[i+2,k,:,:] is the ith varname data at the kth requested timestep, sorted in 
                      the same way as the centroids.

      Note that the data is re-ordered in INCREASING coordinate, i.e. bottom to top in z.

      If return_map is True, then returns a tuple, (data, map) where
      map is a (NX,NZ) array of integers specifying which global id
      corresponds to the (i,j) cell.  This is useful for mapping input
      data back INTO the unstructured mesh.

    Example usage:  
      Calculate and plot the thaw depth at step 5.

      // Pull saturation ice -- TD is where sat ice = 0."
      data = transect_data(['saturation_ice', 5)

      // x coordinate for plotting
      x = data[0,0,:,0]
    
      // for each column, find highest z where sat_ice > 0.
      td_i = np.array([np.where(data[2,0,i,:] > 0.)[0][-1] for i in range(data.shape[2])])

      // now that we have an index into the highest cell with ice, determine td as the 
      // mean of the highest cell with ice and the one above that.  Note this assumes
      // all columns have some thawing.
      td_z = np.array( [  (dat[1,0,i,td_i[i]] + dat[1,0,i,td_i[i+1]]) / 2. 
                              for i in range(len(td_i)) ] )

      plt.plot(x, td_z)

    """
    if coord_order is None:
        coord_order = ['x','z']

    if type(varnames) is str:
        varnames = [varnames,]

    # get centroids
    xyz = mesh.meshElemCentroids(mesh_filename, directory)

    # round to avoid issues
    xyz = np.round(xyz, decimals=5)

    # get ordering of centroids
    dtype = [(coord_order[0], float), (coord_order[1], float)]
    num_order = []
    for i in coord_order:
        if i == 'x':
            num_order.append(0)
        elif i == 'y':
            num_order.append(1)
        elif i == 'z':
            num_order.append(2)

    xyz_sort_order = np.array([tuple([xyz[i,x] for x in num_order]) for i in range(len(xyz))], dtype=dtype)
    xyz_sorting = xyz_sort_order.argsort(order=coord_order)

    with h5py.File(os.path.join(directory,filename),'r') as dat:
        keys_avail = dat[fullname(varnames[0])].keys()
        #keys_avail.sort(lambda a,b: int.__cmp__(int(a),int(b)))
        sorted(keys_avail)

        if keys == 'all':
            keys = keys_avail
        elif type(keys) is str:
            keys = [keys,]
        elif type(keys) is int:
            keys = [keys_avail[keys],]
        elif type(keys) is slice:
            keys = keys_avail[keys]
        elif type(keys) is list:
            if all(type(k) is int for k in keys):
                keys = [keys_avail[k] for k in keys]
            elif all(type(k) is str for k in keys):
                pass
            else:
                raise RuntimeError("Keys requested cannot be processed -- should be 'all', int, or str key, or list of ints or strs.")
                

        # get data
        vals = np.zeros((len(varnames)+2, len(keys), len(xyz)), 'd')

        for i,key in enumerate(keys):
            if deformable:
                xyz = mesh.meshElemCentroids(mesh_filename, directory)
            vals[0,i,:] = xyz[xyz_sorting,num_order[0]]
            vals[1,i,:] = xyz[xyz_sorting,num_order[1]]
            for j,varname in enumerate(varnames):
                vals[j+2,i,:] = dat[fullname(varname)][key][:,0][xyz_sorting]

    # reshape the data
    # determine nx
    nx = len(set(vals[0,0,:]))
    nz = vals.shape[2] / nx
    if (nx * nz != vals.shape[2]):
        raise RuntimeError("Assumption about first coordinate being cleanly binnable is falling apart -- ask Ethan to rethink this algorithm!")
    shp = vals.shape

    if not return_map:
        return vals.reshape(int(shp[0]), int(shp[1]), int(nx), int(nz))
    else:
        return vals.reshape(int(shp[0]), int(shp[1]), int(nx), int(nz)), xyz_sorting.reshape(int(nx), int(nz))


def plot(dataset, ax, cax=None, vmin=None, vmax=None, cmap="jet",
         label=None, mesh_filename="ats_vis_mesh.h5", directory=".", y_coord=0.0,
         linewidths=1):
    """Draws a dataset on an ax."""
    import matplotlib.collections
    from matplotlib import pyplot as plt

    if vmin is None:
        vmin = dataset.min()
    if vmax is None:
        vmax = dataset.max()

    # get the mesh and collapse to 2D
    etype, coords, conn = mesh.meshElemXYZ(filename=mesh_filename, directory=directory)
    if etype != 'HEX':
        raise RuntimeError("Only works for Hexs")

    coords2 = np.array([[coords[i][0::2] for i in c[1:] if abs(coords[i][1] - y_coord) < 1.e-8] for c in conn])
    try:
        assert coords2.shape[2] == 2
        assert coords2.shape[1] == 4
    except AssertionError:
        print(coords2.shape)
        for c in conn:
            if len(c) != 9:
                print(c)
                raise RuntimeError("what is a conn?")
            coords3 = np.array([coords[i][:] for i in c[1:] if abs(coords[i][1] - y_coord) < 1.e-8])
            if coords3.shape[0] != 4:
                print(coords)
                raise RuntimeError("Unable to squash to 2D")

    # reorder anti-clockwise
    for i,c in enumerate(coords2):
        centroid = c.mean(axis=0)
        def angle(p1,p2):
            a1 = np.arctan2((p1[1]-centroid[1]),(p1[0]-centroid[0]))
            a2 = np.arctan2((p2[1]-centroid[1]),(p2[0]-centroid[0]))
            if a1 < a2:
                return -1
            elif a2 < a1:
                return 1
            else:
                return 0

        c2 = np.array(sorted(c,key=functools.cmp_to_key(angle)))
        coords2[i] = c2

    polygons = matplotlib.collections.PolyCollection(coords2, edgecolor='k', cmap=cmap, linewidths=linewidths)
    polygons.set_array(dataset)
    polygons.set_clim(vmin,vmax)
    ax.add_collection(polygons)

    xmin = min(c[0] for c in coords.values())
    xmax = max(c[0] for c in coords.values())
    zmin = min(c[2] for c in coords.values())
    zmax = max(c[2] for c in coords.values())

    ax.set_xlim(xmin,xmax)
    ax.set_ylim(zmin,zmax)

    if cax != None:
        cb = plt.colorbar(polygons, cax=cax)
        if label != None:
            cb.set_label(label)
        
    return ((xmin,xmax),(zmin,zmax))
    
