#!/usr/bin/env python3

import matplotlib.colors as colors
import matplotlib.pyplot as plt
import numpy as np
import math
import os

with open('./topography.txt', "r") as data:
    # while True:
    #     p = data.tell()
    #     line = data.readline()
    #     col1, col2, col3 = line.strip().split()
    #     if col2 == 'ncols': NXCELL = int(col3) #number of cells in x direction
    #     if col2 == 'nrows': NYCELL = int(col3) #number of cells in y direction
    #     if not line.startswith('#'):
    #         data.seek(p) # go back one line
    #         break

    xtp, ytp, ztp = np.loadtxt(data, delimiter=' ', unpack=True)

xmax   = (xtp[0]+xtp[-1])                   # domain extent in x-direction
ymax   = (ytp[0]+ytp[-1])                   # domain extent in y-direction
NXCELL = int(math.sqrt(len(xtp)*xmax/ymax)) # number of cells in x-direction
NYCELL = int(len(xtp)/NXCELL)               # number of cells in y-direction

# first reshape to 2-D array then rotate by ninety degrees and flip in
# vertical direction to conform to FullSWOF indexing convention
x_co = np.flipud(np.rot90(np.reshape(xtp, (NXCELL,NYCELL)))) # x co-ordinates
y_co = np.flipud(np.rot90(np.reshape(ytp, (NXCELL,NYCELL)))) # y co-ordinates
elev = np.flipud(np.rot90(np.reshape(ztp, (NXCELL,NYCELL)))) # elevation map

# calculate cell size in x and y directions
DX = x_co[0,1] - x_co[0,0]
DY = y_co[1,0] - y_co[0,0]
#print(DX)

#print(y_co[0])
#plt.plot(x,z, label='elevation')
#plt.xlabel('x / m')
#plt.ylabel('z / m')

# initialize boundary cells
bb = np.zeros((1, NXCELL))     # bottom boundary
tb = np.zeros((1, NXCELL))     # top boundary
lb = np.zeros((1, NYCELL + 2)) # left boundary
rb = np.zeros((1, NYCELL + 2)) # right boundary
# use forward/backward differencing to calculate boundary elevation values
for i in range(0, NXCELL):
    bb[0,i] = 2.0*elev[0,i] - elev[1,i]
    tb[0,i] = 2.0*elev[NYCELL-1,i] - elev[NYCELL-2,i]

# add boundary cells
z2 = np.concatenate((np.concatenate((bb, elev)), tb))
z3 = np.concatenate((np.concatenate((lb.T, z2), axis=1), rb.T), axis=1)
#print(z2[NYCELL + 1])
#print(z3.shape)

# initialize gradient dz/dy
dzdy = np.zeros((NYCELL, NXCELL))
# use central differencing to calculate nodal values
for i in range(1, NYCELL + 1):
    for j in range(1, NXCELL + 1):
        dzdy[i-1, j-1] = (z3[i+1, j] - z3[i-1, j])/(2.0*DY)
# Note: dz/dy index and co-ordinates are related by:
#            index = int(co-ord/DX - 0.5)
#print("dzdy", dzdy[NYCELL-1, 280:301])
#print("elev", elev[NYCELL-1, 280:301])

xMarker = []
yMarker = []

xch = [] # /pixels
ych = [] # /pixels
zch = [] # /metres

def detach_display():
    fig, (ax1,ax2) = plt.subplots(1, 2, sharey='row')
    #fig, ax1 = plt.subplots()

    left = ax1.imshow(elev,
                      interpolation='nearest',
                      origin='lower', 
                      cmap=plt.get_cmap('terrain_r'))
    ax1.set_title('elevation / m')
    # ax1.set_xticks([0, 100, 200, 300, 400, 500])
    ax1.plot(xMarker, yMarker, 'ro')
    fig.colorbar(left, ax=ax1)

    right = ax2.imshow(dzdy,
                       interpolation='nearest',
                       origin='lower',
                       cmap=plt.get_cmap('rainbow'))
                       # norm=colors.Normalize(vmin=0, vmax=0.1))
    ax2.set_title('slope ($dz/dy$)')
    # ax2.set_xticks([0, 100, 200, 300, 400, 500])
    fig.colorbar(right, ax=ax2)

    # press a suitable key (one that is not bound already) to mark bank.
    def on_key(event):
        if event.inaxes is not None:
            # print(np.rint(event.xdata).astype(int),
            #       np.rint(event.ydata).astype(int))
            xMarker.append(np.rint(event.xdata).astype(int))
            yMarker.append(np.rint(event.ydata).astype(int))
            ax1.plot(event.xdata, event.ydata, 'ro')
            fig.canvas.draw()
        else:
            print('Key pressed ouside axes bounds but inside plot window')
            # TODO: check xMarker, yMarker have even number of items
            # (all left bank and right bank markers are present).


            #print(elev.min())
            # Note: row-index is plotted vertically, column-index plotted
            # horizontally.
            #print(np.unravel_index(elev.argmin(), elev.shape))

    cid = fig.canvas.callbacks.connect('key_press_event', on_key)
    #plt.figure()
    plt.show()
    fig.canvas.callbacks.disconnect(cid)
    #print(xMarker)

def plot_curve(PX):
    """
    Fit straight line to channel elevation data.  Pixel size PX is
    used to convert pixel indices to co-ordinate values.
    """
    # sort markers by y value
    #yx = list(zip(yMarker, xMarker))
    #yx.sort()
    #x_sorted = [x for y, x in yx]
    #y_sorted = [y for y, x in yx]
    #print(x_sorted)
    #print(y_sorted)

    xch.clear()
    ych.clear()
    zch.clear()
    for i in range(0, len(xMarker), 2):
        length = int(np.hypot(xMarker[i+1]-xMarker[i],
                              yMarker[i+1]-yMarker[i]))
        #print(length)
        x = np.linspace(xMarker[i+1], xMarker[i], length)
        y = np.linspace(yMarker[i+1], yMarker[i], length)
        # Extract the values along the line.  Note index order: row,
        # column.
        zi = elev[y.astype(np.intp), x.astype(np.intp)]
        # identify minimum value
        zch.append(zi.min())
        ind = np.unravel_index(zi.argmin(), zi.shape)
        xch.append(x[ind].astype(np.intp))
        ych.append(y[ind].astype(np.intp))
        #print('xch = ', xch[i//2], ' ych = ', ych[i//2], ' zch = ', zch[i//2])

    # TODO: Define starting point in channel.  Find nearest neighbour to
    # starting point.  Extract from xch, ych, zch.  Create new list with
    # extracted point.  Find nearest neighbour to extracted point.  Extract.
    # Add to list.  Repeat.

    # Fit with polyfit
    m, c = np.polyfit([(0.5 + y)*PX for y in ych], zch, 1)
    print('gradient =', m, 'intercept =', c)

    fig, ax = plt.subplots()
    line1, = ax.plot([(0.5 + y)*PX for y in ych], zch)
    line2, = ax.plot([(0.5 + y)*PX for y in ych],
                     [c + m*(0.5 + y)*PX for y in ych], '--')
    ax.set_xlabel('Distance / m')
    ax.set_ylabel('Height / m')
    ax.set_title('Channel profile')
    #plt.figure()
    plt.show()
    #fig.clear()
    #plt.clf()
    #plt.cla()
    #plt.close()
    #fig.canvas.draw()
    #fig.canvas.flush_events()

def save_xyz(PX):
    """
    Write channel elevation data to file.  Pixel size PX is used to
    convert pixel indices to co-ordinate values.
    """
    with open('1D.txt', 'w') as f:
        for xitem, yitem, zitem in zip(reversed(xch),
                                       reversed(ych),
                                       reversed(zch)):
            f.write('{:.2f} {:.2f} {:.4f}\n'.format((0.5 + xitem)*PX,
                                                    (0.5 + yitem)*PX,
                                                    zitem))
    

EXIT_COMMAND = "q"

# if os.fork():
#     # parent
#     pass
# else:
#     # child
#     detach_display()

#detach_display()

while (True):
    input_str = input('Locate markers (m), plot channel profile (p), save profile (s) or exit (q): ')
    #print("input_str = {}".format(input_str))

    if (input_str == EXIT_COMMAND):
        print("End.")
        break
    elif(input_str == "m"):
        detach_display()
    elif(input_str == "p"):
        plot_curve(DY)
    elif(input_str == "s"):
        save_xyz(DY)
    
#print("End.")