import numpy as np
import matplotlib.pyplot as plt
import volume_frac as vfrac


def find_t_x(func = vfrac.variable_growth_rate, dt=0.1, flat_change = 0.0000001,
        r = 0, N = 1.6 * 10 ** 14, T = 300, M0 = 1e4, Q = 156000, P_D = 1e5):
    t_list = [0]
    x_list = [0]

    t_list.append(t_list[-1] + dt)
    x_list.append(func(t = t_list[-1], r = r, N = N, T = T, M0 = M0, Q = Q, P_D=P_D))

    while x_list[-1] < 0.5: #starts finding x-values
        t_list.append(t_list[-1] + dt)
        x_list.append(func(t_list[-1], r = r, N = N, T = T, M0 = M0, Q = Q, P_D=P_D))

    while abs(x_list[-1] - x_list[-2]) > flat_change: #continiues until the graph flats out
        t_list.append(t_list[-1] + dt)
        x_list.append(func(t_list[-1], r = r, N = N, T = T, M0 = M0, Q = Q, P_D=P_D))

    return np.array(t_list), np.array(x_list)

def plot_S_curve(func = vfrac.variable_growth_rate, dt = 0.001, flat_change = 0.0000001, label = 'S-curve',
                 r=0, N=1.6 * 10 ** 14, T=300, M0=1e4, Q=156000, P_D=1e5):
    t_list, x_list = find_t_x(func = func, dt = dt, flat_change = flat_change,
                              r=r, N= N, T=T, M0=M0, Q=Q, P_D=P_D)

    if r != 0:
        label = 'Variable growth rate, r = ' + str(r)

    plt.plot(t_list, x_list, label = label)
    plt.xscale('log')
    plt.xlabel('Time [s]')
    plt.ylabel('Volume fraction recrystalized [–]')


def find_avrami(func = vfrac.variable_growth_rate, dt = 0.001, flat_change = 0.0000001, get_coeff = False,
                r=0, N=1.6 * 10 ** 14, T=300, M0=1e4, Q=156000, P_D=1e5):

    t_list, x_list = find_t_x(func = func, dt = dt, flat_change = flat_change,
                              r=r, N=N, T=T, M0=M0, Q=Q, P_D=P_D)

    #omit the first values to avoid division by zero
    avrami_y = np.log(t_list[2:])
    avrami_x = np.log(np.log(1/(1-x_list[2:])))

    if get_coeff == True:
        coeff = np.polyfit(avrami_x, avrami_y, 1)
        return coeff

    return avrami_x, avrami_y

def plot_avrami(func = vfrac.variable_growth_rate, dt = 0.001, flat_change = 0.0000001, label = 'Avrami plot',
                r=0, N=1.6 * 10 ** 14, T=300, M0=1e4, Q=156000, P_D=1e5):

    avrami_x, avrami_y = find_avrami(func = func, dt = dt, flat_change = flat_change,
                                        r=r, N=N, T=T, M0=M0, Q=Q, P_D=P_D)
    if r != 0:
        label = 'Variable growth rate, r = ' + str(r)

    plt.title('Avrami-plot')
    plt.plot(avrami_x, avrami_y, label = label)
    plt.xscale('linear')
    plt.xlabel(r'$\ln\left(\ln\left[\frac{1}{1-X_V}\right]\right)$')
    plt.ylabel(r'$\ln(K) + n\ln(t)$')