import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 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


def track_path(args):
    H = np.array([[1, 1, 0, 0],
                  [0, 1, 1, 0],
                  [0, 0, 1, 1]])

    homotopy = homotopy_generator.HomotopyGenerator(H)

    print(f"G: {homotopy.G}")
    print(f"F: {homotopy.F}")
    print(f"H: {homotopy.H}")
    print(f"DH: {homotopy.DH}")

    tracker = path_tracker.PathTracker(homotopy, 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 = [], [], [], []

    y = np.zeros(5) + np.array([0.5, 0.48, -0.1, 0.2, 0])
    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}")

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

    for y in np.transpose(ys):
        plt.plot(y)

    plt.show()


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("-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()