import numpy as np
import matplotlib.pyplot as plt


### Properties ###
steps = 100

# Data
T_data = np.array([265, 255, 245, 235, 225, 215])
il_data = np.array([75, 60, 45, 30, 15, 0])
albedo_data = np.array([0.15, 0.25, 0.35, 0.45, 0.55, 0.65])

# Ice latitude
m1, b1 = np.polyfit(T_data, il_data, 1)

def calc_ice_latitude(T):
    return m1 * T + b1

# Planetary albedo
m2, b2 = np.polyfit(T_data, albedo_data, 1)

def calc_albedo(T):
    return np.clip(m2 * T + b2, 0, 1)

### Constants and Data ###
epsilon = 1
sigma = 5.67E-8                         # W/m2 K4

### Energy balance equation ###
def calc_temperature(albedo, L):
    return pow((L * (1 - albedo)) / (4 * epsilon * sigma), 1/4)

### Iteration step ###
def iterate(L, initial_albedo, steps):
    T = calc_temperature(initial_albedo, L)
    for i in range(steps):
        albedo = calc_albedo(T)
        new_T = calc_temperature(albedo, L)
        T = new_T
    return T, albedo

### Loop ###
l_range_bounds = [ 1200, 1600 ]


current_albedo = 0.15
temp_down = []
l_range_down = []
L = l_range_bounds[1]
while L > l_range_bounds[0]-1:
    T, current_albedo = iterate(L, current_albedo, steps)
    temp_down.append(T)
    L = L - 10
    l_range_down.append(L)

current_albedo = 0.65
temp_up = []
l_range_up = []
L = l_range_bounds[0]
while L < l_range_bounds[1]+1:
    T, current_albedo = iterate(L, current_albedo, steps)
    temp_up.append(T)
    L = L + 10
    l_range_up.append(L)


### Input ###
# L, albedo, nIters = input("").split()
# L, albedo, nIters = [ float(L), float(albedo), int(nIters) ]

# T, albedo = iterate(L, albedo, nIters)
# print(T, " ", albedo)

### Plot hysteresis ###
plt.plot(l_range_down, temp_down, label="Cooling")
plt.plot(l_range_up, temp_up, label="Warming")
plt.xlabel("Solar Constant (L)")
plt.ylabel("Equilibrium Temperature (K)")
plt.legend()
plt.show()