# Copyright (c) 2025 Piotr Grabowski
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>

import numpy

from integrator import *
import matplotlib.animation as animation
import matplotlib; matplotlib.use("TkAgg")

import scipy.signal as signal
import scipy.stats as stats
c = 300000
G = 6.67430e-2
m1 = 1.989e6
m2 = 0.3285

mu = m1 * m2 / (m1 + m2)
M = m1 + m2


def NewtonianPotential(r):
    return - G * m1 * m2 / r



def CarthesianFromRelative(q):
    Q_rel = q
    return [m2/M * Q_rel, -m1/M * Q_rel]


class NewtonianHamiltonian:
    @staticmethod
    def H(q, p):
        P_rel = p
        Q_rel = q

        return np.linalg.norm(P_rel)**2 / (2 * mu) + NewtonianPotential(np.linalg.norm(Q_rel))

    @staticmethod
    def dHdp(q, p):
        return p/mu


    @staticmethod
    def dHdq(q, p):
        return np.array([
            G * m1 * m2 * q[0]/np.linalg.norm(q)**3,
            G * m1 * m2 * q[1]/np.linalg.norm(q)**3,
            G * m1 * m2 * q[2]/np.linalg.norm(q)**3
        ])


def l(x):
    return np.linalg.norm(x)

class PostNewtonianHamiltonian:
    @staticmethod
    def H(q, p):
        P_rel = p
        Q_rel = q

        P1 = m2/M * P_rel
        P2 = -m1/M * P_rel

        R = l(Q_rel)

        return np.linalg.norm(P_rel)**2 / (2 * mu) + NewtonianPotential(R) \
            + 1/c**2 * (1/8 * (l(P1) ** 4/m1**3 + l(P2)**4/m2**3) + G**2 * (m1 + m2)*m1*m2/(2*R**2) + G * m1 * m2/(8 * R) * (-12 * l(P1)**2/m1**2 - 12 * l(P2)**2/m2**2 - 28 * l(P_rel)**2/M**2 - 4/M**2 * np.dot(Q_rel/R, P_rel)**2))

    @staticmethod
    def dHdp(q, p):
        P_len = l(p)
        r = l(q)
        n = q/r

        return np.array([
            p[0]/mu - 1/(8 * M**4 * c**2) * 2 * P_len**2 * 2 * p[0] * (m2**4/m1**3 + m1**4/m2**3) - G*m1*m2/(8*c**2*r) * (2 * p[0] * 1/M**2 * (12 * (m1**2/m2**2 + m2**2/m1**2) + 28) + 4/M**2 * np.dot(n, p) * n[0]),
            p[1]/mu - 1/(8 * M**4 * c**2) * 2 * P_len**2 * 2 * p[1] * (m2**4/m1**3 + m1**4/m2**3) - G*m1*m2/(8*c**2*r) * (2 * p[1] * 1/M**2 * (12 * (m1**2/m2**2 + m2**2/m1**2) + 28) + 4/M**2 * np.dot(n, p) * n[1]),
            p[2]/mu - 1/(8 * M**4 * c**2) * 2 * P_len**2 * 2 * p[2] * (m2**4/m1**3 + m1**4/m2**3) - G*m1*m2/(8*c**2*r) * (2 * p[2] * 1/M**2 * (12 * (m1**2/m2**2 + m2**2/m1**2) + 28) + 4/M**2 * np.dot(n, p) * n[2])
        ])

    @staticmethod
    def dHdq(q, p):
        P_len = l(p)
        r = l(q)
        n = q/r

        n_dot_p = np.dot(n, p)

        return np.array([
            G * m1 * m2 * q[0]/r**3 - G**2 * M * m1 * m2 * q[0]/(2 * c**2 * r**4) + G*m1*m2*q[0]/(8 * c**2 * r ** 3) * (P_len**2/M**2 * (12 * (m2**2/m1**2 + m1**2/m2**2) + 28) + 4/M**2 * n_dot_p**2) - G * m1 * m2/(c**2 * r * M**2) * n_dot_p * p[0] * (r**2 - q[0]**2)/r**3,
            G * m1 * m2 * q[1]/r**3 - G**2 * M * m1 * m2 * q[1]/(2 * c**2 * r**4) + G*m1*m2*q[1]/(8 * c**2 * r ** 3) * (P_len**2/M**2 * (12 * (m2**2/m1**2 + m1**2/m2**2) + 28) + 4/M**2 * n_dot_p**2) - G * m1 * m2/(c**2 * r * M**2) * n_dot_p * p[1] * (r**2 - q[1]**2)/r**3,
            G * m1 * m2 * q[2]/r**3 - G**2 * M * m1 * m2 * q[2]/(2 * c**2 * r**4) + G*m1*m2*q[2]/(8 * c**2 * r ** 3) * (P_len**2/M**2 * (12 * (m2**2/m1**2 + m1**2/m2**2) + 28) + 4/M**2 * n_dot_p**2) - G * m1 * m2/(c**2 * r * M**2) * n_dot_p * p[2] * (r**2 - q[2]**2)/r**3,
        ])

def main():
    hamiltonian = PostNewtonianHamiltonian
    ORBIT = 1/1e6 * 88 * 24 * 60 * 60
    AU = 1.496e2
    dt = ORBIT / 25000
    speedup = 1
    q, p, t = integrate(hamiltonian=hamiltonian, p0=np.array([0.0, 5.897e1 * mu, 0.0]), q0=np.array([46.00, 0.0, 0.0]), t=500 * ORBIT, dt0=dt, omega=4/dt, n0=4)
    
    np.savetxt("q.txt", q)
    np.savetxt("p.txt", p)
    np.savetxt("t.txt", t)

    # q = np.loadtxt("q.txt")
    # p = np.loadtxt("p.txt")
    # t = np.loadtxt("t.txt")

    def energy(i):
        return hamiltonian.H(q[i], p[i])

    energies = np.array([energy(i) for i in range(len(q))])
    energies = (energies - energies[0])/energies[0]
    energies = energies**2
    plt.plot(t/ORBIT * 88/365, energies)
    plt.xlabel("Czas [yrs]")
    plt.ylabel("Względny błąd energii")
    plt.title("Względny błąd energii (SI4)")
    plt.show()

    v = numpy.trapezoid(energies, t)
    print("Integral of energy dev: ", v)
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    qs = np.array([CarthesianFromRelative(qi) for qi in q])/AU
    ax.plot(qs[:, 0, 0], qs[:, 0, 1], qs[:, 0, 2], 'r', label='Ciało 1')
    ax.plot(qs[:, 1, 0], qs[:, 1, 1], qs[:, 1, 2], 'b', label='Ciało 2')
    plt.xlabel("Położenie względem środka masy [AU]")
    plt.title("Orbity ciał (SI4)")
    plt.legend();
    plt.show()

    qlens = np.array([l(qi) for qi in q])
    plt.plot(t/ORBIT * 88/365, qlens, label='Simulation Data')
    plt.xlabel("t [yrs]")
    plt.ylabel("Distance to the Sun")

    minima = signal.argrelextrema(qlens, np.less, order=3)[0]
    plt.plot((t[minima])/ORBIT * 88/365, qlens[minima], '.r', label='Minima')
    plt.legend();
    plt.show()

    angleAtMinimum = np.array([(np.atan2(q[minima[i]][1], q[minima[i]][0]) * 180 / np.pi) for i in range(len(minima))])
    arcsecAtMinimum = angleAtMinimum * 3600
    xSpace = np.array([i+1 for i in range(len(arcsecAtMinimum))]) * 88/365
    plt.plot(xSpace, arcsecAtMinimum)
    plt.title("Precesja peryhelium Merkurego (SI4)")
    plt.xlabel("Czas [yrs]")
    plt.ylabel("Położenie kątowe peryhelium [arcsec]")

    slope, intercept, r_value, p_value, std_err = stats.linregress(xSpace, arcsecAtMinimum)
    plt.plot(xSpace, [slope * xValue + intercept for xValue in xSpace], color='r')

    print(slope)

    plt.show()


if __name__ == "__main__":
    main()