import matplotlib.pyplot as plt
import numpy as np

def dT_dt(t, T, r = 0.1, T_s = 22):
    return -r*(T - T_s)

#Tar inn x_n, y_n, dy/dx og steglengde
#returnerer y_{n+1} etter ett rk4-steg
def rk4_step_func(x, y, h, dy_dx):
    k1 = dy_dx(x,y)
    k2 = dy_dx(x + h/2, y + h* k1 / 2)
    k3 = dy_dx(x + h/2, y + h* k2 / 2)
    k4 = dy_dx(x + h, y + h*k3)

    return y + (h/6) * (k1 + 2*k2 + 2*k3 + k4)

#Tar inn en funksjon som utfører ett numerisk steg,
#Starttemperatur og numerisk steglengde
#Returnerer lister for temperaturutviklingen dersom man kjøler 5 grader med en gang og
#hvis man kjøler 5 grader når koppen når 70 grader
def optimum_cooling_time(num_step_func = rk4_step_func, T_0 = 90, dt = 0.1, T_cold = 65, cream_cooling = 5):
    late_T = T_0
    late_t = 0

    late_T_list = [late_T]
    late_t_list = [late_t]

    while late_T > T_cold + cream_cooling:
        late_T = num_step_func(late_t, late_T, dt, dT_dt)
        late_T_list.append(late_T)

        late_t += dt
        late_t_list.append(late_t)
    late_T -= cream_cooling

    late_T_list.append(late_T)
    late_t_list.append(late_t)


    early_T = T_0 - cream_cooling
    early_t = 0
    early_T_list = [T_0, early_T]
    early_t_list = [0, 0]

    while early_T > T_cold:
        early_T = num_step_func(early_t, early_T, dt, dT_dt)
        early_T_list.append(early_T)

        early_t += dt
        early_t_list.append(early_t)

    return early_t_list, early_T_list, late_t_list, late_T_list


#Bruker optimum_cooling_time for å plotte temputviklingen i de to tilfellene beskrevet der
def plot_optimum_cooling(T_0 = 80, dt = 0.02, T_cold = 65, cream_cooling = 5):

    early_t_list, early_T_list, late_t_list, late_T_list = \
        optimum_cooling_time(T_0 = T_0, dt = dt, T_cold = T_cold, cream_cooling=cream_cooling)

    plt.plot(early_t_list, early_T_list, color='blue', label='Early cream')
    plt.plot(late_t_list, late_T_list, color='red', label='Late cream')

    plt.xlabel('Time [min]')
    plt.ylabel('Temp [°C]')

    plt.legend()