import argparse
import numpy as np
import pandas as pd

# autopep8: off
import sys
import os
sys.path.append(f"{os.path.dirname(os.path.abspath(__file__))}/../")

# TODO: How do I import PathTracker and HomotopyGenerator properly?
from hccd import path_tracker, homotopy_generator
# autopep8: on


class ToyHomotopy:
    """Helper type implementing necessary functions for PathTracker.

       Toy example homotopy:
           G = [[x1],
                [x2]]

           F = [[x1 + x2 ],
                [x2 + 0.5]]

           H = (1-t)*G + t*F

       Note that
           y := [[x1],
                 [x2],
                 [t]]
    """
    @staticmethod
    def evaluate_H(y: np.ndarray) -> np.ndarray:
        """Evaluate H at y."""
        x1 = y[0]
        x2 = y[1]
        t = y[2]

        result = np.zeros(shape=2)
        result[0] = x1 + t * x2
        result[1] = x2 + t * 0.5

        return result

    @staticmethod
    def evaluate_DH(y: np.ndarray) -> np.ndarray:
        """Evaluate Jacobian of H at y."""
        x1 = y[0]
        x2 = y[1]
        t = y[2]

        result = np.zeros(shape=(2, 3))
        result[0, 0] = 1
        result[0, 1] = t
        result[0, 2] = x2
        result[1, 0] = 0
        result[1, 1] = 1
        result[1, 2] = 0.5

        return result


def track_path(args):
    tracker = path_tracker.PathTracker(ToyHomotopy, args.euler_step_size, args.euler_max_tries,
                                       args.newton_max_iter, args.newton_convergence_threshold, args.sigma)

    ys_start, ys_prime, ys_hat_e, ys = [], [], [], []

    try:
        y = np.zeros(3)
        for i in range(args.num_iterations):
            y_start, y_prime, y_hat_e, y = tracker.transparent_step(y)
            ys_start.append(y_start)
            ys_prime.append(y_prime)
            ys_hat_e.append(y_hat_e)
            ys.append(y)
            print(f"Iteration {i}: {y}")
    except Exception as e:
        print(f"Error: {e}")

    ys_start = np.array(ys_start)
    ys_prime = np.array(ys_prime)
    ys_hat_e = np.array(ys_hat_e)
    ys = np.array(ys)

    df = pd.DataFrame({"x1b": ys_start[:, 0],
                       "x2b": ys_start[:, 1],
                       "tb":  ys_start[:, 2],
                       "x1p": ys_prime[:, 0],
                       "x2p": ys_prime[:, 1],
                       "tp":  ys_prime[:, 2],
                       "x1e": ys_hat_e[:, 0],
                       "x2e": ys_hat_e[:, 1],
                       "te":  ys_hat_e[:, 2],
                       "x1n": ys[:, 0],
                       "x2n": ys[:, 1],
                       "tn":  ys[:, 2]
                       })

    if args.output:
        df.to_csv(args.output, index=False)
    else:
        print(df)


def main():
    parser = argparse.ArgumentParser(
        description='Homotopy continuation path tracker')

    parser.add_argument("--verbose", default=False, action='store_true')
    parser.add_argument("--euler-step-size", type=float,
                        default=0.05, help="Step size for Euler predictor")
    parser.add_argument("--euler-max-tries", type=int, default=5,
                        help="Maximum number of tries for Euler predictor")
    parser.add_argument("--newton-max-iter", type=int, default=5,
                        help="Maximum number of iterations for Newton corrector")
    parser.add_argument("--newton-convergence-threshold", type=float,
                        default=0.01, help="Convergence threshold for Newton corrector")
    parser.add_argument("-s", "--sigma", type=int, default=1,
                        help="Direction in which the path is traced")
    parser.add_argument("-o", "--output", type=str, help="Output csv file")
    parser.add_argument("-n", "--num-iterations", type=int, default=20,
                        help="Number of iterations of the example program to run")

    args = parser.parse_args()

    track_path(args)


if __name__ == '__main__':
    main()