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


def plot_2d(df: pd.DataFrame):
    df = pd.read_csv("temp.csv")

    fig = plt.figure()
    plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95)
    ax = fig.add_subplot()

    # Analytically computed solution
    ts = df['tb']
    x2s = -ts / 2
    x1s = -ts * x2s

    ax.scatter(x1s, x2s)

    ax.set_xlabel("x1")
    ax.set_ylabel("x2")
    # ax.set_zlabel("t")

    df = df.reset_index()

    for _, row in df.iterrows():
        x1b = row['x1b']
        x2b = row['x2b']
        # x1p = row['x1p']
        # x2p = row['x2p']
        x1e = row['x1e']
        x2e = row['x2e']
        x1n = row['x1n']
        x2n = row['x2n']

        ax.plot([x1b, x1e], [x2b, x2e], 'b')
        ax.plot([x1e, x1n], [x2e, x2n], 'r')

    plt.show()


def plot_3d(df: pd.DataFrame):
    df = pd.read_csv("temp.csv")

    fig = plt.figure()
    plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95)
    ax = fig.add_subplot(projection='3d')

    # Analytically computed solution
    ts = df['tb']
    x2s = -ts / 2
    x1s = -ts * x2s

    ax.scatter(x1s, x2s)
    ax.scatter(x1s, x2s, ts)

    ax.set_xlabel("x1")
    ax.set_ylabel("x2")
    ax.set_zlabel("t")

    df = df.reset_index()

    for _, row in df.iterrows():
        x1b = row['x1b']
        x2b = row['x2b']
        tb = row['tb']
        # x1p = row['x1p']
        # x2p = row['x2p']
        # tp  = row['tp']
        x1e = row['x1e']
        x2e = row['x2e']
        te = row['te']
        x1n = row['x1n']
        x2n = row['x2n']
        tn = row['tn']

        ax.plot([x1b, x1e], [x2b, x2e], 'b', zs=[tb, te])
        ax.plot([x1e, x1n], [x2e, x2n], 'r', zs=[te, tn])

    plt.show()


def main():
    # Parse command line arguments

    parser = argparse.ArgumentParser()

    parser.add_argument("--input", "-i", type=str, required=True,
                        default="temp.csv", help="Filename of the csv file")
    parser.add_argument("--projection", "-p", type=str, required=False,
                        default="2d", help="2d or 3d plot")

    args = parser.parse_args()

    # Plot data
    df = pd.read_csv(args.input)

    if args.projection == "2d":
        plot_2d(df)
    elif args.projection == "3d":
        plot_3d(df)


if __name__ == "__main__":
    main()