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 RepetitionCodeHomotopy:
#     """Helper type implementing necessary functions for PathTracker.
#
#        Repetiton code homotopy:
#            G = [[x1],
#                 [x2],
#                 [x1]]
#
#            F = [[1 - x1**2],
#                 [1 - x2**2],
#                 [1 - x1*x2]]
#
#            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]
#
#         print(y)
#
#         result = np.zeros(shape=3)
#         result[0] = -t*x1**2 + x1*(1-t) + t
#         result[1] = -t*x2**2 + x2*(1-t) + t
#         result[2] = -t*x1*x2 + x1*(1-t) + t
#
#         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=(3, 3))
#         result[0, 0] = -2*t*x1 + (1-t)
#         result[0, 1] = 0
#         result[0, 2] = -x1**2 - x1 + 1
#         result[1, 0] = 0
#         result[1, 1] = -2*t*x2 + (1-t)
#         result[1, 2] = -x2**2 - x2 + 1
#         result[1, 0] = -t*x2 + (1-t)
#         result[1, 1] = -t*x1
#         result[1, 2] = -x1*x2 - x1 + 1
#
#         return result


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

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