#!/usr/bin/env python3

import sys
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import axes3d, Axes3D

plt.rcParams['text.usetex'] = True

f = open( sys.argv[1], 'r' )
current_sigma = 0.0
current_rho = 0.0
current_beta = 0.0
sigma_lst = []
rho_lst = []
beta_lst = []
x_lst = []
y_lst = []
z_lst = []
x_data = []
y_data = []
z_data = []
parameters = []
data = {}
size = 0
for line in f:
   line = line.strip()
   a = line.split()
   if not a:
      continue
   if a[ 0 ] == "#":
      if x_lst:
         parameters_tuple = ( current_sigma, current_rho, current_beta )
         parameters.append( parameters_tuple )
         data_tuple = ( np.array( x_lst ), np.array( y_lst ), np.array( z_lst ) )
         data[ parameters_tuple ] = data_tuple
         x_lst.clear()
         y_lst.clear()
         z_lst.clear()

      current_sigma = float( a[ 2 ] )
      current_rho = float( a[ 4 ] )
      current_beta = float( a [ 6 ] )
      if current_sigma not in sigma_lst:
         sigma_lst.append( current_sigma )
      if current_rho not in rho_lst:
         rho_lst.append( current_rho )
      if current_beta not in beta_lst:
         beta_lst.append( current_beta )
   else:
      x_lst.append( float( a[ 0 ] ) )
      y_lst.append( float( a[ 1 ] ) )
      z_lst.append( float( a[ 2 ] ) )
parameters_tuple = ( current_sigma, current_rho, current_beta )
parameters.append( parameters_tuple )
data_tuple = ( np.array( x_lst ), np.array( y_lst ), np.array( z_lst ) )
data[ parameters_tuple ] = data_tuple

sigma_n = len( sigma_lst )
sigma_idx = 1
for sigma in sigma_lst:
   rho_n   = len( rho_lst )
   beta_n  = len( beta_lst )

   fig, ax = plt.subplots( rho_n, beta_n, figsize=(8, 8), sharey=True, sharex=True )
   fig.suptitle( f'$\sigma={sigma}$')
   #ax = Axes3D(fig)  dos not work with ax indexing
   rho_idx = 0
   beta_idx = 0
   for rho in rho_lst:
      for beta in beta_lst:
         parameters_tuple = ( sigma, rho, beta )
         data_tuple = data[ parameters_tuple ]
         ax[ rho_idx, beta_idx ].plot( data_tuple[ 1 ], data_tuple[ 2 ], linewidth=1.0 )
         if beta_idx == 0:
            ax[ rho_idx, beta_idx ].set_ylabel( f'$\\rho={rho}$' )
         if rho_idx == rho_n-1:            
            ax[ rho_idx, beta_idx ].set_xlabel( f'$\\beta={beta}$' )
         beta_idx = beta_idx + 1
      beta_idx = 0
      rho_idx = rho_idx + 1

   plt.savefig( f"{sys.argv[2]}-{sigma_idx}.png" )
   sigma_idx = sigma_idx + 1
   plt.close(fig)