import numpy as np
from scipy.integrate import quad
from scipy.optimize import fsolve
import matplotlib.pyplot as plt

def integrand_q(z, q, x, y):
  """Integrand for the q equation."""
  return np.exp(-0.5 * z**2) * np.tanh( (q**0.5 * z + np.exp(y)) / x )**2

def integrand_x(z, q, x, y):
  """Integrand for the x equation."""
  return np.exp(-0.5 * z**2) * 1 / np.cosh( (q**0.5 * z + np.exp(y)) / x )**4

def solve_for_y(x):
  """Solves for y given x."""

  def equations(vars):
    """Equations to be solved."""
    q, y = vars
    
    # Calculate q
    q_val, _ = quad(integrand_q, -np.inf, np.inf, args=(q, x, y))
    q_val /= np.sqrt(2 * np.pi)

    # Calculate x^2
    x2_val, _ = quad(integrand_x, -np.inf, np.inf, args=(q, x, y))
    x2_val /= np.sqrt(2 * np.pi)

    return [q - q_val, x**2 - x2_val]

  # Initial guess for q and y
  q_guess = 1.0
  log_y_guess = 0.0

  # Solve the system of equations
  sol = fsolve(equations, (q_guess, log_y_guess))
  
  return np.exp(sol[1])
xs = np.linspace(0.01, 0.95, 500)
ys = [solve_for_y(x) for x in xs]

nmin, nmax = 10, 460
plt.plot(xs[nmin:nmax], ys[nmin:nmax], color = 'black')

# plot y = sqrt(4/3 (1-x)**3)
xs_high = np.linspace(xs[nmax], 1, 200)
plt.plot(xs_high, np.sqrt(4/3 * (1-xs_high)**3), color = 'black')

# plot x = 4/(3 sqrt(2 pi)) * e^(-y^2/2)
ys_high = np.linspace(ys[nmin], 5, 200)
plt.plot(4/(3 * np.sqrt(2*np.pi)) * np.exp(-ys_high**2/2), ys_high, color = 'black')

plt.xlabel('x')
plt.ylabel('y')
plt.title('de Almeida-Thouless line')
plt.grid(True)
plt.savefig("de_Almeida_Thouless_line.svg")
plt.show()