#!/usr/bin/env python3

import matplotlib.pyplot as plt
import numpy as np
import bisect
import ast
import math

#TODO: use YAML/ruamel.yaml for configuration file.
def read_definition(filename):
    ddict = {}
    with open(filename, "r") as f:
        for line in f:
            items = line.split(': ', 1)
            if len(items) == 2:
                ddict[items[0]] = ast.literal_eval(items[1])
    return ddict

# read boundary definition file.
definition_dict = read_definition('boundaryDefinition.txt')
#for dd in definition_dict:
#    print(definition_dict[dd])
slope       = abs(definition_dict["slope"])  # slope at top boundary
target_flow = definition_dict["target_flow"] # imposed discharge
location    = definition_dict["location"]    # boundary location
plotting    = definition_dict["plotting"]    # enable or disable plotting
n_co        = definition_dict["n_co"]        # Manning's 'n' coefficients
# TODO: use weighted mean 'n' values.  See
# http://help.floodmodeller.com/isis/ISIS/River_Section.htm (Eq. 4)
# Note: weighted mean calculation requires roughness map.
height_data = definition_dict["height_data"] # topography
markers     = definition_dict["markers"]     # distances from corner point
panel       = definition_dict["panel"]       # panel fill order
ztol        = definition_dict["ztol"]        # tolerance in overtopping height
numH        = definition_dict["numH"]        # number of height intervals


# print(len(markers))

# with open('./1D_top.txt', "r") as data:
#     xch, ych, zch = np.loadtxt(data, delimiter=' ', unpack=True)

# Fit with polyfit
# m, c = np.polyfit(ych, zch, 1)
# print('gradient =', m, 'intercept =', c)

with open(height_data, "r") as topo:
    xtp, ytp, ztp = np.loadtxt(topo, delimiter=' ', unpack=True)

# domain extent in x-direction:
xmax = (xtp[0]+xtp[-1])
# domain extent in y-direction:
ymax = (ytp[0]+ytp[-1])
# number of cells in x-direction:
ncols = int(math.sqrt(len(xtp)*xmax/ymax))
# number of cells in y-direction:
nrows = int(len(xtp)/ncols)
# cell size
dX = xmax/ncols
print('dX =', dX)

#print(ncols, nrows)

# extract slices from height data array.  Note: xyz format uses ncols
# blocks, with nrows lines per block.
if location == 'top':
    xin = xtp[nrows-1:len(xtp):nrows]
    yin = ytp[nrows-1:len(xtp):nrows]
    zin = ztp[nrows-1:len(xtp):nrows]
    outputFilename = "BCTop.txt"
elif location == 'bottom':
    xin = xtp[0:len(xtp):nrows]
    yin = ytp[0:len(xtp):nrows]
    zin = ztp[0:len(xtp):nrows]
    outputFilename = "BCBottom.txt"
elif location == 'left':
    xin = xtp[:nrows]
    yin = ytp[:nrows]
    zin = ztp[:nrows]
    outputFilename = "BCLeft.txt"
elif location == 'right':
    xin = xtp[nrows*(ncols-1):]
    yin = ytp[nrows*(ncols-1):]
    zin = ztp[nrows*(ncols-1):]
    outputFilename = "BCRight.txt"
    

print(xin)

num_panels = len(panel) # number of panels across boundary

# convert marker co-ordinates to array indices:
marker_ind = [0]
for i in range(len(markers)):
    marker_ind.append(int(markers[i]/dX - 1/2))
if location == 'left' or location == 'right':
    marker_ind.append(nrows-1)
elif location == 'top' or location == 'bottom':
    marker_ind.append(ncols-1)


# print(marker_ind)

xregion = []
zregion = []
for p in range(num_panels):
    xregion.append(xin[marker_ind[p]:marker_ind[p+1]])
    zregion.append(zin[marker_ind[p]:marker_ind[p+1]])

# xregion_west = xin[100:281]
# zregion_west = zin[100:281]

# xregion_east = xin[300:408]
# zregion_east = zin[300:408]

# print(zregion)
# print(xin[12:20])

zmin = []
for p in range(num_panels):
    # identify minimum heights within each panel:
    zmin.append(zregion[p].min())

print(zmin[1])

# channel overtopping height (minimum of left bank and right bank heights):
zmax = min(zregion[panel[0]][0], zregion[panel[0]][-1]) - ztol

def conveyance(numH, n_co, xregion, zregion, zmin, zmax):
    p_i = []                # wetted perimeter
    A_i = []                # area
    r_h = []                # hydraulic radius
    h_i = []                # list of heights
    K_i = []                # conveyance
    Q_i = []                # discharge
    x_sub = [[] for i in range(numH)] # list of x values in subregion
    z_sub = [[] for i in range(numH)] # list of z values in subregion
    for i in range(numH):
        h_i.append(zmin + (i+1)*(zmax-zmin)/numH)
        #print(zregion[zregion < h_i[i]])
        booleanArray = zregion < h_i[i]
        #print(booleanArray[i])
        x_sub[i] += list(xregion[booleanArray])
        z_sub[i] += list(zregion[booleanArray])
        for interval in range(len(xregion)-1):
            if booleanArray[interval+1] != booleanArray[interval]:
                x_extra = xregion[interval] + (h_i[i]-zregion[interval])*(xregion[interval+1]-xregion[interval])/(zregion[interval+1]-zregion[interval])
                bisect.insort(x_sub[i], x_extra) # add intercept value
                ind_x = x_sub[i].index(x_extra)
                z_sub[i].insert(ind_x, h_i[i])   # add height value
        #print(z_sub[i])

        dp = 0
        dA = 0
        eps = 1e-06
        for j in range(len(x_sub[i])-1):
            if abs(z_sub[i][j+1] - h_i[i]) > eps or abs(z_sub[i][j] - h_i[i]) > eps:
                dp += np.hypot(x_sub[i][j+1] - x_sub[i][j],
                               abs(z_sub[i][j+1] - z_sub[i][j]))
            #print(dp)
            # calculate area using trapezium rule
            dA += (h_i[i] - (z_sub[i][j+1] + z_sub[i][j])/2)*(x_sub[i][j+1] - x_sub[i][j])
            #print('Area =', dA)

        p_i.append(dp)
        A_i.append(dA)

        r_h.append(A_i[i]/p_i[i])  # ratio of area and wetted perimeter
        #print('hydraulic radius =', r_h[i])
        K_i.append(A_i[i]*(1/n_co)*r_h[i]**(2/3)) # conveyance
        Q_i.append(K_i[i]*slope**0.5)             # discharge

    return p_i, A_i, r_h, h_i, K_i, Q_i

#print(h_i)

def plot_region(xdata, labelx,
                ydata1, labely1,
                ydata2, labely2,
                ydata3, labely3, titlep):
    plt.xlabel(labelx)
    plt.ylabel(labely1)
    plt.title(titlep)
    plt.plot(xdata, ydata1)
    plt.show()

    plt.xlabel(labelx)
    plt.ylabel(labely2)
    plt.title(titlep)
    plt.plot(xdata, ydata2)
    plt.show()

    plt.xlabel(labelx)
    plt.ylabel(labely3)
    plt.title(titlep)
    plt.plot(xdata, ydata3)
    plt.show()

p_i = [[] for _ in range(num_panels)]
A_i = [[] for _ in range(num_panels)]
r_h = [[] for _ in range(num_panels)]
h_i = [[] for _ in range(num_panels)]
K_i = [[] for _ in range(num_panels)]
Q_i = [[] for _ in range(num_panels)]
for p in range(num_panels):
    p_i[p], A_i[p], r_h[p], h_i[p], K_i[p], Q_i[p] = conveyance(
        numH,
        n_co[p],
        xregion[p],
        zregion[p],
        zmin[p],
        zmax)
    if plotting:
        plot_region(
            h_i[p]-zmin[p], 'maximum depth / m',
            r_h[p], 'hydraulic radius / m',
            K_i[p], r'conveyance / $m^3/s$',
            Q_i[p], r'discharge / $m^3/s$',
            'Panel {}'.format(p))

# sort list of discharge lists according to panel fill order:
sortedQ = [Q_i[i] for i in panel]
# create cumulative discharge list:
total_flow = np.cumsum([item[-1] for item in sortedQ])
print(total_flow)
# target_flow_west = target_flow - Q_i[-1] - Q_i_east[-1]
# calculate velocity: note dependence on hydraulic radius
velocity_channel = Q_i[panel[0]][-1]/A_i[panel[0]][-1]
# velocity_east    = Q_i_east[-1]/A_i_east[-1]

# print(target_flow_west)
# find part-filled panel:
ind_p = bisect.bisect(total_flow, target_flow)
print(ind_p)

# calculate target flow in part-filled panel:
panel_target_flow = target_flow - total_flow[ind_p-1]

# find insertion point for target flow value:
ind_q = bisect.bisect(Q_i[panel[ind_p]], panel_target_flow) 
print(ind_q)

# find height at target flow by linear interpolation
h_extra = h_i[panel[ind_p]][ind_q-1] + (h_i[panel[ind_p]][ind_q]-h_i[panel[ind_p]][ind_q-1])*(panel_target_flow-Q_i[panel[ind_p]][ind_q-1])/(Q_i[panel[ind_p]][ind_q]-Q_i[panel[ind_p]][ind_q-1])
print(h_i[panel[ind_p]][ind_q-1], h_extra, h_i[panel[ind_p]][ind_q])
# find area at target flow by linear interpolation
A_extra = A_i[panel[ind_p]][ind_q-1] + (h_extra-h_i[panel[ind_p]][ind_q-1])*(A_i[panel[ind_p]][ind_q]-A_i[panel[ind_p]][ind_q-1])/(h_i[panel[ind_p]][ind_q]-h_i[panel[ind_p]][ind_q-1])
print(r_h[panel[ind_p]][ind_q-1], r_h[panel[ind_p]][ind_q])

velocity_panel    = panel_target_flow/A_extra
print(velocity_channel, velocity_panel)

csa = np.zeros(len(xin))             # cross-sectional area of element
csa_p = np.zeros(num_panels)         # cross-sectional area of panel
for i, p in enumerate(panel):
    if i < ind_p:
        # panels are filled
        area_sum = 0
        for m in range(marker_ind[p], marker_ind[p+1]):
            csa[m] = max(0, (zmax - zin[m])*dX)
            area_sum += csa[m]
        csa_p[p] = area_sum
    elif i == ind_p:
        # panel is part-filled
        area_sum = 0
        for m in range(marker_ind[p], marker_ind[p+1]):
            csa[m] = max(0, (h_extra - zin[m])*dX)
            area_sum += csa[m]
        csa_p[p] = area_sum
    else:
        # panel is empty
        csa_p[p] = 0


#print('csa_p[0] = {} csa_p[1] = {}'.format(csa_p[0], csa_p[1]))
#print('A_i_west = {} A_i = {} A_i_east = {}'.format(A_extra, A_i[-1], A_i_east[-1]))

def save_bc():
    with open(outputFilename, 'w') as f:
        f.write('{:6} {:>2} {:>2} {:>10}\n'.format('#x', 'c', 'q', 'h'))
        for ind_z, (xitem, zitem) in enumerate(zip(xin, zin)):
            panel_x = bisect.bisect(markers, xitem)
            if csa[ind_z] == 0:
                # wall boundary condition
                f.write('{:6.2f} {:>2}\n'.format(xitem, 2))
            elif panel_x == ind_p:
                # imposed discharge within part-filled panel
                f.write('{:6.2f} {:>2} {:10.6f} {:9.6f}\n'.format(
                    xitem, 5,
                    -csa[ind_z]*panel_target_flow/csa_p[panel_x],
                    h_extra-zitem))
            else:
                # imposed discharge within filled panels
                f.write('{:6.2f} {:>2} {:10.6f} {:9.6f}\n'.format(
                    xitem, 5,
                    -csa[ind_z]*Q_i[panel_x][-1]/csa_p[panel_x],
                    zmax-zitem))

save_bc()