import numpy
import scipy.stats as stats
import scipy.optimize as optimize
import scipy.interpolate as interpolate


solids_concentration = stats.norm(loc=62.8, scale=3.14)
mass_flow_rate = stats.norm(loc=1250, scale=150)


Beach_Slopes = """	0.39	0.38	0.36	0.35	0.33	0.3	0.27
56.65	0.58	0.6	0.61	0.62	0.64	0.67	0.71
59.19	0.73	0.75	0.77	0.78	0.81	0.84	0.9
61.15	0.87	0.89	0.91	0.93	0.96	1.01	1.07
62.8	1.02	1.05	1.07	1.09	1.13	1.18	1.25
64.45	1.2	1.24	1.26	1.29	1.33	1.39	1.48
66.41	1.49	1.54	1.57	1.6	1.66	1.73	1.84
68.95	2.03	2.08	2.13	2.17	2.25	2.36	2.51"""

slope_rows = Beach_Slopes.split('\n')
slope_x_indices = map(float, slope_rows[0].split('\t')[1:])
slope_y_indices = map(float, [r.split('\t')[0] for r in slope_rows[1:]])
slope_data = [map(float, r.split('\t')[1:]) for r in slope_rows[1:]]

slope_interpolator = interpolate.interp2d(slope_x_indices, slope_y_indices, slope_data)
    
# copy/pasted from previous tutorial       
terrain_length = 2000

weights = numpy.ones((terrain_length,))
weights[0] = 0.5
weights[-1] = 0.5

base_grade = 0.01  # 1% slope 

ground_array = numpy.zeros((terrain_length,))
x_coordinates = numpy.linspace(0,terrain_length-1,terrain_length)
gradient = x_coordinates * base_grade
ground_array = ground_array - gradient

from muk3d.geometry.polyline import PolyLine
from muk3d.view import add_to_scene

polyline = PolyLine()
polyline.new_line_from_points(zip(x_coordinates,  [0,] * terrain_length,  ground_array))

add_to_scene('tailings profile', polyline, overwrite_existing=True)

def pour_tailings(discharge_elevation, beach_slope, current_ground):
    tailings_surface = discharge_elevation - x_coordinates * beach_slope
    
    new_tailings_surface = numpy.maximum(current_ground, tailings_surface)
    area = numpy.sum((new_tailings_surface- current_ground) * weights)
    
    return area, new_tailings_surface

current_ground = ground_array.copy()
    
def optimise_z(discharge_elevation, beach_slope, current_ground, target_area):
    area, new_ground = pour_tailings(discharge_elevation, beach_slope, current_ground)
    return (area - target_area)/target_area

target_area = 100

deposited_area = 0
for i in range(100):
    z_guess = max(current_ground) + 1
    slope = slope_interpolator(mass_flow_rate.rvs(), solids_concentration.rvs())/100.0
    final_z = optimize.newton(optimise_z, z_guess, tol=0.001, args=(slope, current_ground, target_area))

    deposition_2d_area, tailings_surface = pour_tailings(final_z, slope, current_ground)

    polyline.new_line_from_points(zip(x_coordinates,  [0,] * terrain_length,  tailings_surface))
    deposited_area += deposition_2d_area
    current_ground = tailings_surface
    
print 'deposited area', deposited_area