import matplotlib.pyplot as plt
import numpy as np
import argparse
from copy import deepcopy
import os,sys

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(SCRIPT_DIR))

from formatlist import assignformat

def addline(x, y, form:dict, label=None, formatindicator="line-dot"):
 if formatindicator == "dot":
 plt.scatter(x,y,c=form["c"], s=10, marker=form["marker"], label=label)
 elif formatindicator == "line-dot":
 plt.plot(x,y,c=form["c"], linestyle=form["linestyle"], marker=form["marker"], label=label)
 elif formatindicator == "line":
 plt.plot(x,y,c=form["c"], linestyle=form["linestyle"], label=label)


def setfigform(xtickList, ytickList, xlabel, ylabel, title = ""):
 plt.subplots_adjust(left=0.2, right=0.9, top=0.9, bottom=0.15)
 plt.legend()
 plt.title(title,fontsize = 16)
 plt.xlabel(xlabel,fontsize = 14)
 plt.ylabel(ylabel, fontsize = 14)
 plt.xticks( xtickList)
 plt.yticks( ytickList)
 plt.tick_params(axis='both',width=2,labelsize = 12)
 

def getmaxmin(data, dtype = float):
 tmpdata = np.copy(data)
 max_data = np.reshape(tmpdata, [-1]).max()
 min_data = np.reshape(tmpdata, [-1]).min()
 return max_data, min_data

def readcontext(context, num_y=1, skip_y=0):
 x = []
 y = []
 for i in range(num_y):
 y.append([])
 for line in context:
 c = line.split()
 if len(c) < num_y+skip_y+1:
 break
 x.append(float(c[0]))
 for idx in range(skip_y+1, num_y+skip_y+1):
 y[idx -1-skip_y].append(float(c[idx]))
 return x,y

if __name__ == "__main__":
 labellist = ["500K", "600K", "700K", "800K"]

 parser = argparse.ArgumentParser(description='Figure texts')
 parser.add_argument('--title', type=str, default='', help='titile of the figure')
 parser.add_argument('--skip', type=int, default=0, help='skip columes')
 parser.add_argument('--skip_first_line', type=bool, default=True, help="skip first line")
 parser.add_argument('--num_y', type=int, default=1, help="number of y")
 parser.add_argument('--xlabel', type=str, default="x", help="xlabel")
 parser.add_argument('--ylabel', type=str, default="y", help="ylabel")
 parser.add_argument('INPUT', type=str,
 help="input file")
 parser.add_argument('--format', type=str, default='line-dot', help="Format of plots: line, line-dot, dot")
 parser.add_argument("--diagonal_line", type=bool, default=False, help="Add diagonal line")
 args = parser.parse_args()
 
 formatindicator = args.format
 
 inputfile = args.INPUT
 
 num_y = args.num_y
 skip_first_line = args.skip_first_line
 skip_y = args.skip
 
 
 fin = open(inputfile, "r")
 context = fin.readlines()
 if skip_first_line:
 context.pop(0)

 x, y = readcontext(context, num_y = num_y, skip_y = skip_y)
 
 print(x)
 for yi in y:
 print(yi)
 print("\n")
 for i in range(num_y):
 form = assignformat[formatindicator]
 addline(x,y[i],form[i],labellist[i],formatindicator=formatindicator)
 
 
 max_x, min_x = getmaxmin(x)
 max_y, min_y = getmaxmin(y)
 print((min_x, min_y))
 print((max_x, max_y))
 xtickList = (max_x-min_x) * np.arange(0, 1, 0.2) + min_x
 ytickList = (max_y-min_y) * np.arange(0, 1, 0.2) + min_y
 
 setfigform(xtickList, ytickList, xlabel = args.xlabel, ylabel = args.ylabel, title = args.title)
 # add diagonal line
 if args.diagonal_line:
 plt.plot((min_x, max_x), (min_y, max_y), ls="--", c="k")
 
 plt.savefig("fig")
 plt.show()