# 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 sys

import numpy as np
import matplotlib.pyplot as plt


def print_progress(iteration, total, prefix='', suffix='', decimals=1, bar_length=100):
    str_format = "{0:." + str(decimals) + "f}"
    percents = str_format.format(100 * (iteration / float(total)))
    filled_length = int(round(bar_length * iteration / float(total)))
    bar = '█' * filled_length + '-' * (bar_length - filled_length)

    sys.stdout.write('\r%s |%s| %s%s %s' % (prefix, bar, percents, '%', suffix)),

    if iteration == total:
        sys.stdout.write('\n')
    sys.stdout.flush()


def rotMatrix(dt, omega):
    return np.array([[np.cos(2 * omega * dt), np.sin(2 * omega * dt)],
                     [-np.sin(2 * omega * dt), np.cos(2 * omega * dt)]], dtype=np.longdouble)


def gamma(l):
    return 1 / (2 - 2 ** (1 / (l + 1)))


def integrate(*, hamiltonian, p0, q0, t, dt0, omega, lambd=1, n0=2):
    def phiA(q, p, x, y, dt):
        return q, p - dt * hamiltonian.dHdq(q, y), x + dt * hamiltonian.dHdp(q, y), y

    def phiB(q, p, x, y, dt):
        return q + dt * hamiltonian.dHdp(x, p), p, x, y - dt * hamiltonian.dHdq(x, p)

    def phiHc(q, p, x, y, dt):
        # rot = rotMatrix(dt, omega)
        # qp = np.array([q + x, p + y], dtype=np.longdouble) + rot @ np.array([q - x, p - y], dtype=np.longdouble)
        # xy = np.array([q + x, p + y], dtype=np.longdouble) - rot @ np.array([q - x, p - y], dtype=np.longdouble)
        #
        # return qp[0]/np.longdouble(2), qp[1]/np.longdouble(2), xy[0]/np.longdouble(2), xy[1]/np.longdouble(2)

        cosa = np.cos(2 * omega * dt * lambd)
        sina = np.sin(2 * omega * dt * lambd)

        return (q * (1 + cosa) + x * (1 - cosa) + lambd * sina * (p - y)) / 2, \
               (p * (1 + cosa) + y * (1 - cosa) - 1/lambd * sina * (q - x)) / 2, \
               (x * (1 + cosa) + q * (1 - cosa) - lambd * sina * (p - y)) / 2, \
               (y * (1 + cosa) + p * (1 - cosa) + 1/lambd * sina * (q - x)) / 2

    def phi2(q, p, x, y, dt):
        # print("dt", dt)
        # q, p, x, y = phiA(q, p, x, y, dt / 2)
        # print("A", q,p,x,y)
        # q, p, x, y = phiB(q, p, x, y, dt / 2)
        # print("B", q,p,x,y)
        # q, p, x, y = phiHc(q, p, x, y, dt)
        # print("C", q,p,x,y)
        # q, p, x, y = phiB(q, p, x, y, dt / 2)
        # print("B", q,p,x,y)
        # q, p, x, y = phiA(q, p, x, y, dt / 2)
        # print("A", q,p,x,y)
        #
        # print('-------')
        # return q,p,x,y
        return phiA(*phiB(*phiHc(*phiB(*phiA(q, p, x, y, dt / 2), dt / 2), dt), dt / 2), dt / 2)


    def phi(n, q, p, x, y, dt):
        if n % 2 == 1:
            raise Exception("order of integrator must be even")

        if n == 2:
            return phi2(q, p, x, y, dt)

        g = gamma(n)
        q1, p1, x1, y1 = phi(n - 2, q, p, x, y, g * dt)
        q2, p2, x2, y2 = phi(n - 2, q1, p1, x1, y1, (1 - 2 * g) * dt)
        q3, p3, x3, y3 = phi(n - 2, q2, p2, x2, y2, g * dt)

        return q3, p3, x3, y3

    t0 = 0
    length = int(t / dt0) + 1
    retvalQ = np.zeros((length, np.shape(q0)[0]), dtype=np.longdouble)
    retvalP = np.zeros((length, np.shape(p0)[0]), dtype=np.longdouble)
    retvalT = np.zeros(length, dtype=np.longdouble)

    retvalQ[0] = q0
    retvalP[0] = p0
    retvalT[0] = t0
    currentQ, currentP, currentX, currentY = q0, p0, q0, p0
    index = 0
    print("Starting")
    print_progress(index, length, prefix='Progress', suffix='Complete')
    while index < length - 1:
        currentQ, currentP, currentX, currentY = phi(n0, currentQ, currentP, currentX, currentY, dt0)
        # print(np.linalg.norm(currentQ - currentX))
        if index % 100 == 0:
            print_progress(index, length, prefix='Progress', suffix='Complete', decimals=3)

        index += 1
        t0 += dt0
        retvalQ[index] = currentQ
        retvalP[index] = currentP
        retvalT[index] = t0

    return np.array(retvalQ), np.array(retvalP), np.array(retvalT)