import numpy
import matplotlib.pyplot


# Inputs
timeDuration = 2000                    # years
timeStep = 10                          # years
waterDepth = 4000                          # meters
initialTemp = 400

# Constants
secondToYear = 60*60*24*365
specificWaterHeatCapacity  = 4.2E3      # kg/m^3
waterDensity = 1000                     # J/kg*m^3
L = 1350                                # Watts/m2
albedo = 0.3
epsilon = 1
sigma = 5.67E-8                         # W/m2 K4


# Calc amount of steps for given duration and resolution
nSteps = int(timeDuration/timeStep)

initialHeat = initialTemp * (waterDepth*waterDensity*specificWaterHeatCapacity)
initialRadiatedHeat = epsilon * sigma * pow(initialTemp,4)
T = [initialTemp]               # K
Q = [initialHeat]               # Watt/m^2
time = [0]                      # years
heatOut = [initialRadiatedHeat] # Watt/m^2
# print(initialRadiatedHeat)

for i in range(nSteps):
    # Calculate the heat change rate 
    dQ_dt = (L * (1 - albedo)/4) - epsilon * sigma * pow(T[-1],4)
    # Add heat change over step period to total heat in system   
    Q.append(Q[-1] + dQ_dt * timeStep * secondToYear)
    # Calculate temperature of planet based on heat in system
    T.append(Q[-1] / (waterDepth*waterDensity*specificWaterHeatCapacity))
    # Calculate only the outgoing heat
    heatOut.append(epsilon * sigma * pow(T[-1],4)) 
    # Save timestep for plotting
    time.append(time[-1] + timeStep)

print(T[-1])


matplotlib.pyplot.plot(time, T)
matplotlib.pyplot.title("Time-dependent temperature of a naked planet")
matplotlib.pyplot.xlabel('Time (years)')
matplotlib.pyplot.ylabel('Temperature (K)')
matplotlib.pyplot.show()