import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np
import matplotlib.patches as patches
import datetime
import os
import sys
from textwrap import wrap

'''if len(sys.argv) > 2:
    os._exit(0)
if len(sys.argv) == 2:
    if os.path.isdir(sys.argv[1]):
        print("Creating png Image.")
    else:
        print("Invalid output directory path.")
        os._exit(0)
else:
    print("Error : Please provide output directory path.")
    os._exit(0)'''

data = []
for x in range(21):
    data.append([])
    for y in range(26):
        if x<3 and y==0:
            data[x].append(10)
        else:
            data[x].append(1)
# create discrete colormap
cmap = colors.ListedColormap(['#00ff00', '#ff0000'])
bounds = [0,10,20]
norm = colors.BoundaryNorm(bounds, cmap.N)

fig, ax = plt.subplots()
ax.imshow(data, cmap=cmap, norm=norm)

# draw gridlines
ax.grid(which='major', axis='both', linestyle='-', color='k')
ax.set_xticks(np.arange(-0.5, 26, 1))
ax.set_yticks(np.arange(-0.5, 21, 1))
labels = [item.get_text() for item in ax.get_xticklabels()]
indx=1
val=0.5
for x in range(26):
  labels[indx] = '%.2f'%(val)
  val=val+0.02
  indx=indx+1
ax.set_xticklabels(labels)

labels = [item.get_text() for item in ax.get_yticklabels()]
indx=0
val=80
for y in range(21):
  labels[indx] = val
  val=val-3
  indx=indx+1
ax.set_yticklabels(labels,fontsize=7)

rect = patches.Rectangle((9.5,9.5),5,5,linewidth=2,edgecolor='b',facecolor='none')

ax.add_patch(rect)

rect = patches.Rectangle((11.5,11.5),1,1,linewidth=2,edgecolor='r',facecolor='none')

ax.add_patch(rect)

#aax.set_title("kona_Main.ATPG_GDSC_SHMOO.SHMOO_MBURST_tk_atpg_int_cvpwrapper_gdsc_coll_0101_XMD_NOM_partF3_TSMC_FFF_t30_site2_sl1_121",fontsize=5,x=0.385, y=1)
title=ax.set_title("\n".join(wrap("kona_Main.ATPG_GDSC_SHMOO.SHMOO_MBURST_tk_atpg_int_cvpwrapper_gdsc_coll_0101_XMD_NOM_partF3_TSMC_FFF_M40_site2_sl1_121",70)),fontsize=8,x=0.5)
plt.xticks(rotation=90,fontsize=6)
plt.xlabel('X-vddcx_1[v](Steps-26,Step Width-0.002)', fontsize=7)
plt.ylabel('Y:Frequency1[MHz](Steps-21,Step Width-3)', fontsize=7)
plt.suptitle("shmoo",x=0.09, y=.95,fontsize=8)
dtStr=datetime.datetime.now()
plt.text(-5,23.25,'\nGenerated by XYZ, on '+dtStr.strftime('%c'),fontsize=5)
plt.savefig("Shmoo_plot.png",dpi=180)
#plt.savefig(sys.argv[1]+"\\Shmoo_plot.png",dpi=180)
#print("Created successfully. Please find png in "+sys.argv[1]+"\\Shmoo_plot.png")
print("Created successfully. Please find png in "+"Shmoo_plot.png")