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
import matplotlib

#fpath = r'C:\Users\User\Desktop\Научная работа\Программы\МОДЭС\test freq\W-35 FISH LOW TEMP 5 TO 60 STEP 3.TXT'
#fpath = r'C:\Users\User\Desktop\Научная работа\Программы\МОДЭС\test freq\P273 NZR -100 TO +200 STEP 5.TXT'
#fpath = r'C:\Users\User\Desktop\Научная работа\Программы\МОДЭС\test freq\NUMBER3 CL 0_03 -150 TO 120 STEP 5.TXT'
fpath = r'D:\Рабочий стол\\МОДЭС\test freq\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") 
            # (из-за ошибки utf-8 encoding)

# 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,2, figsize=(24, 16), dpi = 400,
                         subplot_kw={"projection": "3d"})

# test 3d data

x1 = np.linspace(-5, 5, 10)
y1 = np.linspace(-5, 5, 10)
X1, Y1 = np.meshgrid(x1, y1)
Z1 = np.sin(np.sqrt(X1**2 + Y1**2))

#data from file

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())


ax = axes[0,0]
ax.set_xlabel("Freq")
ax.set_ylabel("Temp")
ax.set_zlabel("Eps")
ax.plot_wireframe(X1, Y1, Z1, cmap= colormap, norm =log_norm)
ax.set_title("3d test data")



ax = axes[0,1]
X, Y = np.meshgrid(x, y)    
Z = np.reshape(z, (y_len, x_len))
ax.set_xlabel("Freq")
ax.set_ylabel("Temp")
ax.set_zlabel("Eps")
ax.plot_wireframe(X, Y, Z, cmap= colormap, norm =log_norm)
ax.set_title("3d file data")



ax = axes[1,0]
y = np.maximum(y, 0)
ax.set_title("file data y >= 0")
ax.set_xlabel("Freq")
ax.set_ylabel("Temp")
ax.set_zlabel("Eps")
X, Y = np.meshgrid(x, y)    
Z = np.reshape(z, (y_len, x_len))
ax.plot_wireframe(X, Y, Z, cmap= colormap, norm =log_norm)

axes[1,1].remove()