import numpy as np
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from matplotlib.colors import Normalize

t_end = 150
t = np.linspace(0, t_end)

p_dahls_nok = 25 + 3 * np.sin(np.pi * t / 50)
p_v1_nok = 10 + 0.05 * t
p_v2_nok = 15 - 0.25 * np.sqrt(t)
p_v3_nok = 20 - 5 * np.cos(np.pi * t / 30)

p_dahls_hc = 30

p_hc = p_dahls_hc / p_dahls_nok

p_v1_hc = p_v1_nok * p_hc
p_v2_hc = p_v2_nok * p_hc
p_v3_hc = p_v3_nok * p_hc

fig, axs = plt.subplots(3, 1, sharex='all')
ax1, ax2, ax3 = axs
twn1 = ax1.twinx()

ax1.plot(t, p_dahls_nok, color='b')
twn1.plot(t, p_hc, color='r')
ax1.set_ylabel(r'$p_{Dahls}$ [NOK]', color='b')
twn1.set_ylabel(r'$p_{HCC}$ [HCC/NOK]', color='r')

ax2.plot(t, p_v1_nok, color='r')
ax2.plot(t, p_v2_nok, color='g')
ax2.plot(t, p_v3_nok, color='b')
ax2.set_ylabel('pris [NOK]')

ax3.plot(t, p_v1_hc, color='r')
ax3.plot(t, p_v2_hc, color='g')
ax3.plot(t, p_v3_hc, color='b')
ax3.set_ylabel('pris [HCC]')

ax3.set_xlabel('t')
plt.show()

p_dahls_nok = lambda t: 25 + 3 * np.sin(np.pi * t / 50)
p_dahls_hc = 30
p_hc = lambda t: p_dahls_hc / p_dahls_nok(t)

t_kjop_hc = [0, 25, 75, 110]

norm = Normalize(vmin=min(t_kjop_hc), vmax=max(t_kjop_hc))
cmap = get_cmap('viridis')

fig, axs = plt.subplots(2, 1, sharex='all')
ax1, ax2= axs
twn1 = ax1.twinx()

ax1.plot(t, p_dahls_nok(t), color='b')
twn1.plot(t, p_hc(t), color='r')
ax1.set_ylabel(r'$p_{Dahls}$ [NOK]', color='b')
twn1.set_ylabel(r'$p_{HCC}$ [HCC/NOK]', color='r')

for ti in t_kjop_hc:
    t_list = np.linspace(ti, t_end)
    ax2.plot(t_list, [p_dahls_nok(tk) - p_dahls_hc / p_hc(ti) for tk in t_list], color=cmap(norm(ti)), label=int(ti))
    ax2.plot([ti], [p_dahls_nok(ti) - p_dahls_hc / p_hc(ti)], marker='x', color=cmap(norm(ti)))
    twn1.plot([ti], [p_hc(ti)], marker='x', color=cmap(norm(ti)))

ax2.set_ylabel('Penger spart [NOK/Dahls]')
ax2.set_xlabel('t')
plt.figlegend(title=r'$t_{kjøp}$', loc='lower right')
plt.show()