import numpy as np
from scipy.optimize import curve_fit
from matplotlib import pyplot as plt
from pyscript import web, when, display

from fipy import Variable, FaceVariable, CellVariable, Grid1D, ExplicitDiffusionTerm, TransientTerm, DiffusionTerm, Viewer
from fipy.tools import numerix

nx = 50
dx = 1.
mesh = Grid1D(nx=nx, dx=dx)

phi = CellVariable(name="solution variable",
                   mesh=mesh,
                   value=0.)


D = 1.

valueLeft = 1
valueRight = 0

phi.constrain(valueRight, mesh.facesRight)
phi.constrain(valueLeft, mesh.facesLeft)

timeStepDuration = 0.9 * dx**2 / (2 * D)
steps = 100

cells = 5

for step in range(steps):
    for j in range(cells):
        phi_new[j] = phi_old[j] \
          + (D * dt / dx**2) * (phi_old[j+1] - 2 * phi_old[j] + phi_old[j-1])
    time += dt


wind_speed = np.array([1.83333333, 1.43333333, 1.25      , 1.98888889, 1.25      ,
       1.25384615, 2.37142857, 1.4       , 1.32222222, 0.60714286,
       0.38461538, 1.47333333, 0.94375   , 1.33333333, 1.7       ,
       1.23571429, 0.45      , 0.81428571, 0.65714286, 0.32666667,
       0.99130435, 1.875     , 0.36428571, 0.87692308, 1.2375    ,
       0.81428571, 1.21111111, 0.91      , 0.73333333, 0.58333333,
       0.45333333, 1.50769231, 1.025     , 0.92      , 0.33529412,
       0.66470588, 1.00714286, 0.7       , 0.5       , 0.54      ,
       0.48461538, 1.3       , 1.5125    , 0.6       , 0.50714286,
       1.20714286])

Q_meas = np.array([0.00880829, 0.0063877 , 0.00581561, 0.00921021, 0.00760904,
       0.00654181, 0.01113902, 0.01070922, 0.00911057, 0.00625395,
       0.00532118, 0.00501692, 0.00499498, 0.00672258, 0.00937989,
       0.0058073 , 0.00535757, 0.01035616, 0.00572211, 0.00550452,
       0.00361425, 0.01065908, 0.00515242, 0.00547755, 0.01025287,
       0.01025955, 0.00860055, 0.00790549, 0.00636647, 0.00590157,
       0.00485068, 0.00484262, 0.00818899, 0.00501573, 0.00448077,
       0.00483092, 0.00520283, 0.0073193 , 0.00468288, 0.00509364,
       0.00644061, 0.00739513, 0.00922624, 0.00780421, 0.00551596,
       0.00541268])



Aw = 156 * 0.0001 

               
def correlation(wind_speed, a, b ):
    Q = (a + b*wind_speed)*Aw
    return Q



# LIN
popt, pcov = curve_fit(correlation, 
          wind_speed, Q_meas,
          p0=[0.25,  0.07]
          )
a = popt[0]
b = popt[1]
display(a,b)

display("ready to compute")

# import meshplex
# import meshzoo
# import numpy as np
# from scipy.sparse import linalg

# import pyfvm
# from pyfvm.form_language import Boundary, dS, dV, integrate, n_dot_grad


# class Poisson:
#     def apply(self, u):
#         return integrate(lambda x: -n_dot_grad(u(x)), dS) - integrate(lambda x: 1.0, dV)

#     def dirichlet(self, u):
#         return [(lambda x: u(x) - 0.0, Boundary())]


# # Create mesh using meshzoo
# vertices, cells = meshzoo.rectangle_tri(
#     np.linspace(0.0, 2.0, 401), np.linspace(0.0, 1.0, 201)
# )
# mesh = meshplex.Mesh(vertices, cells)

# matrix, rhs = pyfvm.discretize_linear(Poisson(), mesh)

# u = linalg.spsolve(matrix, rhs)


# display("DONE pyfvm")



@when("click", "#run-button")
def plot(event):
    """
    Plot.
    """

  
    
    input_text = web.page["constant_a"]
    a = float(input_text.value)
    
    
    Qcor = correlation(wind_speed, a, 0.2 )

    ####
    # wind_speed
    #
    fig, ax = plt.subplots(1)
    
    Vind = np.argsort(wind_speed)
    ax.plot(wind_speed,1000*Q_meas,"dk" , label="meas")
    ax.plot(wind_speed[Vind], 1000*Qcor[Vind],"-", label="cor a=%.2f" % a)
    
    ax.set_ylabel("Volume rate (l/s)")
    
    ax.grid(1)
    ax.set_xlabel("Wind speed (m/s)")
    ax.legend()
    
    
    
    display(fig)