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
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.pyplot import figure as fig
%matplotlib inline
from matplotlib.widgets import Slider

def RecrystallisationKinetics(T, mode): #This is the function which calculates and plots the recrysallisation kinetics, with temperature and mode/ situation as input
    Avrami = [] #List of Avrami values
    tA = [] #The time for the Avrami values
    XA = 0 #The Avrami plot values
    Pd = 10**5 #Stored deformation energy
    M0 = 10**4 #Mobility of the moving boundary of the recrystallized grains (constant) [m/(s*Pa)]
    Q = 156000 # #Activation energy [J/mol]
    N = 2 * 10**12 #Nucleii per unit volume [1/m^3]
    Ndot = 1.6 * 10**14 #Nucleation frequency [1/(m^3*s)]
    R = 8.314  #The ideal gas constant [J/(mol*K)]
    M = M0 * np.exp(-Q/(R * T)) #The temperature deåendent mobility of the recrystallized grains [m/(s*Pa)]
    G = M * Pd #The growth rate [m/s]
    Xl = [] #List of values for the fraction
    X = 0 #Fraction recrystallized material
    r = 0 #Parameter for the varying growth rate case
    if mode == "Johnson": #For the Johnson-Mehl case
        D = (G / Ndot)**(1 / 4) #The grain size [m]
        k = (1 / 3) * np.pi * Ndot * G**3 #k is a parameter where the fraction of recrysallized grains is on the form : X=1-exp(-kt^n)
        n = 4 #As k, n is a parameter in the Avrami solution
        dt = ((np.log(2) / (k))**(1 / n))/ 500 #dt is the time increment [s]
        t = [] #t is a list of the time values
        i = 0 #i is used for making time values in the while loop by multiplying by dt
        while X < 0.99: #The fraction can't be higher than 1, so we want it to stop before that
            i += 1
            t.append(dt*i) #Makes a list of time values
            X = 1 - np.exp(-(1 / 3) * np.pi * Ndot * G**3 * float(t[i - 1])**4)
            XA=np.log(np.log(1 / (1 - X)))
            tA.append(np.log(t[i - 1]))
            Avrami.append(XA)
            Xl.append(X)

        plt.plot(t, Xl, label = 'Fraction recrystallized material') #Plot of the fraction recrysallized material
        plt.title('Fraction recrystallized material, Johnson-Mehl case')
        plt.xlabel('Time [s]')
        plt.ylabel('Fraction [-]')
        plt.legend(loc = "best")
        plt.show()

        plt.plot(tA, Avrami, label='ln(ln(1/1-X(t))') #The Avrami plot
        plt.title('Avrami plot')
        plt.xlabel('ln(t) [ln(s)]')
        plt.ylabel('ln(ln(1/1-X(t))')
        plt.legend(loc = "best")
        plt.show()

    elif mode=="SiteSatC": #For the constant growth case
        D=(1/N)**(1/3)
        k=(4/3)*np.pi*N*G**3
        n=3
        dt =((np.log(2)/k)**(1/n))/500
        t=[]
        i=0
        while X<0.99:
            i+=1
            t.append(dt*i)
            X=1-np.exp(-(4/3)*np.pi*N*G**3*float(t[i-1])**3)
            XA=np.log(np.log(1/(1-X)))
            Xl.append(X)
            tA.append(np.log(t[i-1]))
            Avrami.append(XA)

        plt.plot(t, Xl, label = 'Fraction recrystallized material')
        plt.title('Fraction recrystallized material, Site saturation with constant growth rate')
        plt.xlabel('Time [s]')
        plt.ylabel('Fraction [-]')
        plt.legend(loc = "best")
        plt.show()

        plt.plot(tA, Avrami, label='ln(ln(1/1-X(t))')
        plt.title('Avrami plot, constant growth rate')
        plt.xlabel('ln(t) [ln(s)]')
        plt.ylabel('ln(ln(1/1-X(t))')
        plt.legend(loc = "best")
        plt.show()

    elif mode=="SiteSatVarying": #For the varying growth rate case
        D=(1/N)**(1/3)
        n=3*(1-r) #Here n depends on r
        k=(4/3)*np.pi*N*G**3*(1/(1-r)**3)
        dt= ((np.log(2)/k)**(1/n))/500
        i=0
        t=[]
        for l in range (6): #We want to print a graph for each r, from 0-0.6
            i=0 #Here we reset the values from last iteration
            t=[]
            Xl=[]
            tA=[]
            Avrami=[]
            X=0
            XA=0
            while X<0.99:
                i+=1
                t.append(dt*i)
                X=1-np.exp((-4 * np.pi * N * G**3 * float(t[i-1]**(3*(1-r))))/(3* (1-r)**3))
                XA=np.log(np.log(1/(1-X)))
                Xl.append(X)
                tA.append(np.log(t[i-1]))
                Avrami.append(XA)
            plt.plot(t, Xl, label = 'Fraction recrystallized material')
            plt.title('Fraction recrystallized material, (r=' + str(round(r,1))+")")
            plt.xlabel('Time [s]')
            plt.ylabel('Fraction [-]')
            plt.legend(loc = "best")
            plt.show()

            plt.plot(tA, Avrami, label='ln(ln(1/1-X(t))')
            plt.title('Avrami plot')
            plt.xlabel('ln(t) [ln(s)]')
            plt.ylabel('ln(ln(1/1-X(t))')
            plt.legend(loc = "best")
            plt.show()

            r+=0.1 #r increases by 0.1 for each iteration
    else:
        print("Unknown mode.") #If the mode isn't recognised, the function will print this

    print("The grain size is "+str(D)+"m") #Printing the calculated grain size

T=float(input("Choose temperature: ")) #Input for temperature [K]

def sub_task_chooser(sub): #Code for making buttons
    x = np.linspace(0,np.pi/2)  # x-range

    # Switch function
    if sub == 1:
        mode="Johnson"
    elif sub == 2:
        mode="SiteSatC"
    elif sub == 3:
        mode="SiteSatVarying"
    RecrystallisationKinetics(T, mode)


# Menu toggle buttons
sub_menu = widgets.ToggleButtons(
    #        {'names':corresponding values,}
    options={'Johnson-Mehl':1, 'Constant growth rate':2, 'Varying growth rate':3},
    value = 1,  #default value
    description='Choose case:',  #name
    button_style='info',  #blue buttons
    )
interact(sub_task_chooser, sub=sub_menu)