import os, sys, re
import math as m
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

# path to test file with data
fpath = r'D:\Desktop\MODES\test NUMBER3 CL 0_03 -150 TO 120 STEP 5.TXT'

fpath = os.path.normpath(fpath) 
fname = os.path.split(fpath)[1]

# data search
pattern = r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?"

headlines = 0
with open(fpath, 'r', encoding='ansi') as (f):
    # head lines
    while not re.match(pattern, f.readline().split()[0]):
        headlines += 1
            
data = pd.read_table(fpath, sep='\t', skiprows=headlines, 
            names=['freq', 'temp', "eps_real", "eps_imag","sigma"],                        
            encoding = "ISO-8859-1") 
            
# str to float
try:
    data.applymap(lambda x: float(x))
except:
    # pandas ver. > 2.1.0, map instead applymap !
    data.map(lambda x: float(x))
    
data.sort_values(by=["temp", "freq"], ascending=[True, False], 
                                                   inplace=True)

#----------------------------------------

fig, axes = plt.subplots(2,3, figsize=(18, 12),subplot_kw={"projection": "3d"})

y = data.temp.unique()
y_len = len(y)

fr = data.loc[data['freq'] == data['freq'][0]]
x_len = (fr).index[1]

x = (data['freq'].values)[0:x_len]

z = data['eps_imag'].values

colormap = 'coolwarm' #,  
log_norm = LogNorm(vmin=0.1, vmax=z.max())

def Plot_3D(ax):
    X, Y = np.meshgrid(x, y)   
    Z = np.reshape(z, (y_len, x_len))
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_zscale('log')
    ax.set_xlabel("Freq")
    ax.set_ylabel("Temp")
    ax.set_zlabel("Eps")
    ax.plot_surface(X, Y, Z, cmap= colormap, norm =log_norm)

axes[0,0].remove()

axes[0,1].set_title("X,Y,Z log scale")
Plot_3D(axes[0,1])

axes[0,2].set_title("y >= 0")
y = np.maximum(y, 0)
Plot_3D(axes[0,2])

axes[1,0].set_title("y and z >= 0")
z = np.maximum(z, 0)
Plot_3D(axes[1,0])

axes[1,1].set_title("y >= 1, z >= 0")
y = np.maximum(y, 1)
Plot_3D(axes[1,1])

axes[1,2].set_title("y and z >= 1")
z = np.maximum(z, 1)
Plot_3D(axes[1,2])