import numpy as np
import matplotlib.pyplot as plt
#import ipywidgets as widgets
import warnings
warnings.filterwarnings('ignore')
from ipywidgets import widgets, Output
from IPython.display import display
from IPython.html.widgets import *
from IPython.display import clear_output
#fig2, ax2=plt.subplots()
N_dot = 1.6 * 10 ** 14
N = 2 * 10 ** 12
P = 10 ** 5
M0 = 10 ** 4  # ms/Pa
Q = 156000
T=input(r'Do you want to set a temperature? If yes, enter temperature in $^\circ$C. If not, press enter to set temperature to 300 (default). ')
T=int(T)
T+=273
R=8.314
#T=int(T)

#G=f(T)
#sub_task_chooser(sub_menu)
#print(sub_menu.options)
#T_slider = widgets.IntSlider(
    #value=298,
    #min = 273,
    #max = 1273,
    #step = 1,
    #description='T[K]',
    #continous_update = True)
def f(T):
    T+=273
#T=interact(f, T=widgets.IntSlider(min=0, max=1000, step=1, value=273))
M = M0 * np.exp(-Q/(R*T))
G = M * P
def master(k,n):
    t_half = (np.log(2)/(k))**(1 / n)
    dt = t_half / 50
    t = []
    p = 0
    tid = dt * n
    t_log = []
    X = []
    x = 0
    element = 1
    A = []
    while x<0.99:
        tid=dt*p
        x = 1 - np.exp((-k* tid ** n) / 3)
        if x<0.99:
            t.append(tid)
            X.append(x)
            if tid!=0:
                element = np.log(tid)
                t_log.append(element)
                a=np.log(np.log(1/(1-x)))
                A.append(a)
        p+=1
    return X,A,t_log,t
def sub_task_chooser(sub):
    #x = np.linspace(0,np.pi/2)  # x-range
    # Switch function
    if sub == 1:
        jm=(np.pi * N_dot * G ** 3)/3
        X, A, t_log, t = master(jm, 4)
        plt.plot(t, X)
        plt.show()
        plt.plot(t_log, A)
        plt.show()
        D = (G / N_dot) ** 0.25
        print('Grain size [m]: '+str(D))
    elif sub == 2:
        site_sat=(4*np.pi * N * G ** 3)/3
        X, A, t_log, t = master(site_sat, 3)
        plt.plot(t, X)
        plt.show()
        plt.plot(t_log, A)
        plt.show()
        D = (1 / N) ** (1 / 3)
        print('Grain size [m]: '+str(D))
    elif sub == 3:
        r=0
        fig, ax1 = plt.subplots()
        fig, ax2 = plt.subplots()
        for i in range(7):
            r = i * 0.1
            new_n = 3 * (1 - r)
            new_k = (4 * np.pi * N * G ** 3) / (new_n ** 3)
            X, A, t_log, t = master(new_k, new_n)
            name = 'r=' + str(round(r, 1))
            ax1.plot(t, X, label=name)
            ax1.legend()
            ax2.plot(t_log, A, label=name)
            ax2.legend(loc="lower right")
        plt.show()
        D = (1 / N) ** (1 / 3)
        print('Grain size [m]: ' + str(D))
    #plt.legend()
    #plt.show()

#def buttonclicked(value):
    #sub_task_chooser(value['new'])
sub_menu = widgets.ToggleButtons(
    #        {'names':corresponding values,}
options={'Johnson-Mehl':1, 'Constant growth':2, 'Varying growth':3},
value = 1,  #default value
description='Model:',  #name
button_style='info',  #blue buttons
)

interact(sub_task_chooser, sub=sub_menu)
#interact(master, T=T_slider)