#!/usr/bin/env python3

from boundState import BoundState

from solveSE import WoodsSaxonPot, CoulombPotential, SpinOrbit_Pot, WS_SurfacePot, SolvingSE
from mpmath import coulombf, coulombg

import numpy as np
from scipy.special import gamma, sph_harm, factorial
import matplotlib.pyplot as plt

from assLegendreP import associated_legendre_array

def SevenPointsSlope(data, n):
  return (-data[n + 3] + 9 * data[n + 2] - 45 * data[n + 1] + 45 * data[n - 1] - 9 * data[n - 2] + data[n - 3]) / 60

def FivePointsSlope(data, n):
  return ( data[n + 2] - 8 * data[n + 1] + 8 * data[n - 1] -  data[n - 2] ) / 12

# from sympy.physics.quantum.cg import CG
# from sympy import S
# def clebsch_gordan(j1, m1, j2, m2, j, m):
#     cg = CG(S(j1), S(m1), S(j2), S(m2), S(j), S(m))
#     result = cg.doit()
#     return np.complex128(result)

def KroneckerDelta(i, j):
  if i == j:
    return 1
  else:
    return 0
  
from clebschGordan import clebsch_gordan

############################################################
class DistortedWave(SolvingSE):
  def __init__(self, target, projectile, ELab):
    super().__init__(target, projectile, ELab)
    self.SetRange(0, 0.1, 300)
    self.CalCMConstants()

    self.ScatMatrix = []
    self.distortedWaveU = []

    self.legendreArray = []

  def SetLJ(self, L, J):
    self.L = L
    self.J = J
    self.dsolu0 = pow(0.1, 2*L+1)

  def CoulombPhaseShift(self, L = None, eta = None):
    if L is None:
      L = self.L
    if eta is None:
      eta = self.eta
    return np.angle(gamma(L+1+1j*eta))

  def CalScatteringMatrix(self, maxL = None, verbose = False):
    if maxL is None:
      maxL = self.maxL

    self.ScatMatrix = []
    self.distortedWaveU = []

    for L in range(0, maxL+1):
      sigma = self.CoulombPhaseShift()

      temp_ScatMatrix = []
      temp_distortedWaveU = []

      for J in np.arange(L-self.S, L + self.S+1, 1):
        if J < 0:
          temp_ScatMatrix.append(0)
          temp_distortedWaveU.append(0)
          continue

        self.SetLJ(L, J)
        self.SolveByRK4()

        r1 = self.rpos[-2]
        f1 = float(coulombf(self.L, self.eta, self.k*r1))
        g1 = float(coulombg(self.L, self.eta, self.k*r1))
        u1 = self.solU[-2]

        r2 = self.rpos[-1]
        f2 = float(coulombf(self.L, self.eta, self.k*r2))
        g2 = float(coulombg(self.L, self.eta, self.k*r2))
        u2 = self.solU[-1]

        det = f2*g1 - f1*g2

        A = (f2*u1 - u2*f1) / det
        B = (u2*g1 - g2*u1) / det

        ScatMatrix = (B + A * 1j)/(B - A * 1j)

        if verbose:
          print(f"{{{L},{J}, {np.real(ScatMatrix):10.6f} +  {np.imag(ScatMatrix):10.6f}I}}")

        temp_ScatMatrix.append(ScatMatrix)

        dwU = np.array(self.solU, dtype=np.complex128)
        dwU *= np.exp(-1j*sigma)/(B-A*1j)
        temp_distortedWaveU.append(dwU)
      
      self.ScatMatrix.append(temp_ScatMatrix)
      self.distortedWaveU.append(temp_distortedWaveU)

    return [self.ScatMatrix, self.distortedWaveU]

  def PrintScatteringMatrix(self):
    for L in range(0, len(self.ScatMatrix)):
      print("{", end="")
      for i in range(0, len(self.ScatMatrix[L])):
        print("{", end="")
        print(f"{L:2d},{L+i-self.S:4.1f}, {np.real(self.ScatMatrix[L][i]):10.6f} +  {np.imag(self.ScatMatrix[L][i]):10.6f}I", end="")
        if i < len(self.ScatMatrix[L])-1 :
          print("}, ", end="")
        else:
          print("}", end="")
      print("},")

  def GetScatteringMatrix(self, L, J):
    return self.ScatMatrix[L][J-L+self.S]

  def GetDistortedWave(self, L, J):
    return self.distortedWaveU[L][J-L+self.S]

  def PlotDistortedWave(self, L, J):
    plt.plot(self.rpos, np.real(self.GetDistortedWave(L, J)), label="Real")
    plt.plot(self.rpos, np.imag(self.GetDistortedWave(L, J)), label="Imaginary")
    plt.legend()
    plt.grid()
    plt.show(block=False)
    input("Press Enter to continue...")

  def PlotScatteringMatrix(self):
    nSpin = int(self.S*2+1)
    fig, axes = plt.subplots(1, nSpin, figsize=(6*nSpin, 4))

    for i in range(0, nSpin):
      sm = []
      l_list = []
      for L in range(0, len(self.ScatMatrix)):
        if i == 0 and L == 0 :
          continue
        l_list.append(L)
        sm.append(self.ScatMatrix[L][i])

      axes[i].plot(l_list, np.real(sm), label="Real", marker='o')
      axes[i].plot(l_list, np.imag(sm), label="Imaginary", marker='x')
      axes[i].legend()
      axes[i].set_xlabel('L')
      axes[i].set_ylabel('Value')
      if self.S*2 % 2 == 0 :
        str = f'{int(i-self.S):+d}'
      else:
        str = f'{int(2*(i-self.S)):+d}/2' 
      axes[i].set_title(f'Real and Imaginary Parts vs L for Spin J = L{str}')
      axes[i].grid()

    plt.tight_layout()
    plt.show(block=False)
    input("Press Enter to continue...")

  def RutherFord(self, theta_deg):
    sin_half_theta = np.sin(np.radians(theta_deg + 1e-20) / 2)
    result = self.eta**2 / (4 * (self.k**2) * (sin_half_theta**4))
    return result

  def CoulombScatterintAmp(self, theta_deg):
    sin_sq = pow(np.sin(np.radians(theta_deg + 1e-20)/2), 2)
    coulPS = self.CoulombPhaseShift(0)
    return - self.eta / (2 * self.k * sin_sq) * np.exp(1j * (2*coulPS - self.eta * np.log(sin_sq)))
  
  def GMatrix1Spin(self, v, v0, l ) -> complex:
    if self.S == 0 :
      return self.ScatMatrix[l][0] - KroneckerDelta(v, v0)
    else:
      Jmin = l - self.S
      Jmax = l + self.S
      value = 0
      for J in  np.arange(Jmin, Jmax + 1, 1):
        index = int(J - Jmin)
        cg1 = clebsch_gordan(l, 0, self.S, v0, J, v0)
        cg2 = clebsch_gordan(l, v0 - v, self.S, v, J, v0)
        value += cg1 * cg2 * self.ScatMatrix[l][index]
      return value - KroneckerDelta(v, v0)

  def CalLegendre(self, theta_deg, maxL = None, maxM = None):
    if maxL is None:
      maxL = self.maxL
    if maxM is None:
      maxM = int(2*self.S)
    self.legendreArray = associated_legendre_array(maxL, maxM, theta_deg)
    return self.legendreArray

  def GetPreCalLegendre(self, L, M):
    if abs (M) <= int(2*self.S):
      return self.legendreArray[L][int(abs(M))]
    else :
      return 0

  def NuclearScatteringAmp(self, v, v0, maxL = None ) -> complex:
    value = 0
    if maxL is None:
      maxL = self.maxL
    for l in range(0, maxL+1):
      if abs(v0-v) > l :
        value += 0
      else:
        coulPS = self.CoulombPhaseShift(l)
        fact = pow(-1, v0-v) * np.sqrt(factorial(l - abs(v0-v))/factorial(l + abs(v0-v)))
        value += (2*l+1) * fact * self.GetPreCalLegendre(l, v0-v) * np.exp(2j * coulPS)* self.GMatrix1Spin(v, v0, l)

    return value / 2j / self.k
  
  def DCSUnpolarized(self, theta_deg, maxL = None):
    value = 0
    self.CalLegendre(theta_deg)
    jaja = self.CoulombScatterintAmp(theta_deg)
    for v in np.arange(-self.S, self.S + 1, 1):
      for v0 in np.arange(-self.S, self.S + 1, 1):
        value += abs( jaja * KroneckerDelta(v, v0) +  self.NuclearScatteringAmp(v, v0, maxL))**2

    value = value / (2 * self.S + 1) 
    return value

  def PlotDCSUnpolarized(self, thetaRange = 180, thetaStepDeg = 0.2, maxL = None, verbose = False):
    theta_values = np.linspace(0, thetaRange, int(thetaRange/thetaStepDeg)+1)

    thetaTick = 30
    if thetaRange < 180:
      thetaTick = 10

    y_values = []
    for theta in theta_values:
      if theta == 0:
        y_values.append(1)
      else:
        y_values.append(self.DCSUnpolarized(theta, maxL)/ self.RutherFord(theta))
        if verbose :
          print(f"{theta:6.2f}, {y_values[-1]:10.6f}")

    plt.figure(figsize=(8, 6))
    # plt.plot(theta_values, y_values, marker='o', linestyle='-', color='blue')    
    plt.plot(theta_values, y_values, linestyle='-', color='blue')
    plt.title("Differential Cross Section (Unpolarized)")
    plt.xlabel("Angle [deg]")
    plt.ylabel("D.C.S / Ruth")
    plt.yscale("log")
    plt.xticks(np.arange(0, thetaRange + 1, thetaTick))
    plt.xlim(0, thetaRange)
    plt.grid()
    plt.show(block=False)
    input("Press Enter to continue...")