import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad
from scipy.optimize import newton
from scipy.signal import argrelextrema

palette = iter(['#9b59b6', '#4c72b0', '#55a868', '#c44e52', '#dbc256'])

# Wartosci krytyczne temperatury, objetosci i cisnienia
# Wartosci te odpowiadaja parametrom krytycznym CO2 dla rownania stanu van der Waalsa:
# (p - a/V^2)(V-b) = RT. Jednostki: [p]=Pa, [V]=m^3/mol oraz [T]=K.
pc = 7.404e6
Vc = 1.28e-4
Tc = 304

def vdw(Tr, Vr):
    """Rownanie stanu van der Waalsa.
   Zwroc zredukowane cisnienie w funkcji zredukowanej temperatury i objetosci molowej
    """

    pr = 8*Tr/(3*Vr-1) - 3/Vr**2
    return pr


def vdw_maxwell(Tr, Vr):
    """Rownanie van der Waalsa z przeprowadzona konstrukcja Maxwella.
       Zwroc zredukowane cisnienie w funkcji zredukowanej temperatury i objetosci molowej, 
       stosujac konstrukcje Maxwella dla obszaru niestabilnosci.
    """

    pr = vdw(Tr, Vr)
    if Tr >= 1:
        # Obszar stabilnosci
        return pr

    # Poczatkowe zgdnieciegdzie znajduje sie linia wykorzystywana do konstrukcji Maxwella:
    # objetosc odpowiadajaca wartosci sredniej pomiedzy minimum i maksimum w zredukowanym cisnieniu, pr.

    iprmin = argrelextrema(pr, np.less)
    iprmax = argrelextrema(pr, np.greater)
    Vr0 = np.mean([Vr[iprmin], Vr[iprmax]])

    def get_Vlims(pr0):
        """Rozwiaz rownanie na izoterme van der Waalsa wzgledem zredukowanej objetosci.
        Zwroc minimalna i maksymalna wartosc zredukowanej objetosci takie ze cisnienie 
        zredukowane wynosi pr0. Wzywanie tej funkcji jest uzasadnione tylko gdy T<Tc, tj.
        ponizej temperatury krytyczniej rownanie to posiada trzy pierwiastki.
        """

        eos = np.poly1d( (3*pr0, -(pr0+8*Tr), 9, -3) )
        roots = eos.r
        roots.sort()
        Vrmin, _, Vrmax = roots
        return Vrmin, Vrmax

    def get_area_difference(Vr0):
        """Zwroc roznice pol powierzchni uzyskiwanych z procedury Maxwella dla gazu van der Waalsa.
        Zwroc roznice pomiedzy powierzchniami wyznaczonymi przez krzywe od Vr0 do Vrmax oraz od Vrmin 
        do Vr0, gdzie zredukowane cisnienie uzyskiwane z rownania van der Waalsa jest takie same w Vrmin,
        Vr0 i Vrmax. Dla konstrukcji Maxwella to roznica powierzchni ma być rowna 0.
        """

        pr0 = vdw(Tr, Vr0)
        Vrmin, Vrmax = get_Vlims(pr0)
        return quad(lambda vr: vdw(Tr, vr) - pr0, Vrmin, Vrmax)[0]

    # Szukamy pierwiastkow korzystajac z metody Newtona
    # Wyznaczamy Vr0 odpowiadajace rownym powierzchniom w konstrukcji Maxwella.
    Vr0 = newton(get_area_difference, Vr0)
    pr0 = vdw(Tr, Vr0)
    Vrmin, Vrmax = get_Vlims(pr0)
    # Ustal cisnienie zredukowane w konstrukcji Maxwella w obszarze niestabilnosci
    # by bylo stale i wynosilo pr0. 
    pr[(Vr >= Vrmin) & (Vr <= Vrmax)] = pr0
    return pr

Vr = np.linspace(0.5, 3, 500)

def plot_pV(T):
    Tr = T / Tc
    c = next(palette)
    ax.plot(Vr, vdw(Tr, Vr), lw=2, alpha=0.3, color=c)
    ax.plot(Vr, vdw_maxwell(Tr, Vr), lw=2, color=c, label='{:.2f}'.format(Tr))

fig, ax = plt.subplots()

for T in range(270, 320, 10):
    plot_pV(T)

ax.set_xlim(0.4, 3)
ax.set_xlabel('Zredukowana objetosc, $\omega$')
ax.set_ylim(0, 1.6)
ax.set_ylabel('Zredukowane cisnienie, $\pi$')
ax.legend(title='Zredukowana temperatura')

plt.show()