import numpy as np
import matplotlib.pyplot as plt
import os

nsondes=100
vmin=0
vmax=1300
do_map=1
do_hist=1

head = np.genfromtxt("../exemple_image.txt",max_rows=1)
im=int(head[0])
jm=int(head[1])

def save(tit):
    plt.savefig(tit, bbox_inches="tight", dpi=360)
    plt.close()

def plot_map(tab,tit=None, fig=None):
  plt.figure()
  print('tab[2]=',tab[2])
  marker_size = 60000 / len(tab[1])
  print('s=',marker_size)
  #plt.scatter(tab[1], tab[0], c=tab[2], vmin=vmin,vmax=vmax,s=marker_size,linewidth=0)
  plt.scatter(tab[1], tab[0], c=tab[2], vmin=vmin,vmax=vmax,s=marker_size)
  plt.colorbar()
  plt.xlabel("x")
  plt.ylabel("y")
  if tit: plt.title(tit)
  plt.xlim(0,im)
  plt.ylim(0,jm)
  plt.gca().set_aspect("equal")
  if fig: save(fig)

os.system('pwd')
tab_pred = np.genfromtxt("valeurs_emulees_%isondes.txt"%nsondes).transpose()
tab_sond = np.genfromtxt("valeurs_sondees_%isondes.txt"%nsondes).transpose()
vrai_flux = np.genfromtxt("../exemple_image.txt",skip_header=1).reshape((im,jm))

if do_map:
  plot_map(tab_pred, tit="émulées %i sondes"%nsondes, 
          fig="carte_emulee_%isondes.png"%nsondes)
  plot_map(tab_sond, tit="sondées %i sondes"%nsondes,
          fig="carte_sondee_%isondes.png"%nsondes)

  plt.figure()
  plt.imshow(vrai_flux, origin="lower", vmin=vmin,vmax=vmax)
  plt.colorbar()
  plt.xlabel("x")
  plt.ylabel("y")
  plt.title("z")
  plt.gca().set_aspect("equal")
  save("carte_z.png")

if do_hist:
  plt.figure()
  bins = np.arange(vmin,vmax,3)
  
  def plot_hist(tab, tit, lw=2):
    plt.hist(tab, density=True, bins=bins, label=tit, histtype="step", lw=lw)
  
  plot_hist(tab_pred[2], "emulated")
  plot_hist(tab_sond[2], "sondes", lw=.5)
  plot_hist(vrai_flux.flatten(), "reference")
  plt.xlabel("Flux [W/m2]")
  plt.ylabel("Density")
  plt.legend()
  save("histogram_%isondes.png"%nsondes)