import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

#======================================================
# ==              SCRIPT PARAMETERS                  ==
#======================================================

# === CSV Parameters ===
CSV_FILENAME      =   r"c:\Users\ranny\Downloads\week_2025_41_data.csv"
CSV_SEPARATOR     =   ","

COL_DIAMETER      =   "Diameter1 [mm]"
COL_VELOCITY      =   "Velocity [m/s]"
COL_TEMPERATURE   =   "Temperature [C]"
COL_TEMP_GROUPED  =   "TempGrouped"  # column added by script

# === Limits for Velocity & Diameter Filtering ===
LIM_VELOCITY_MIN  =   0.2   # m/s
LIM_DIAMETER_MIN  =   0.7   # mm
LIM_DIAMETER_MAX  =   5.0   # mm

# === Data Analysis Parameters ===
TEMP_BIN_SIZE     =   2     # split by 2 degrees
FIT_SAMPLE_COUNT  =   100

# === Figure Parameters ===
FIG_WIDTH         =   9
FIG_HEIGHT        =   6

FIG_DATA_ALPHA    =   0.6
FIG_DATA_SIZE     =   35

FIG_TR_COLOR      =   "black"  # trendline color
FIG_TR_SIZE       =   2

FIG_LEGEND_TITLE  =   "Teplotné Skupiny [°C]"
FIG_LEGEND_SIZE   =   8
FIG_LEGEND_COLS   =   2

FIG_TITLE         =   f"Závislosť Rýchlosť a Priemeru Častíc Zoskupené po {TEMP_BIN_SIZE}°C Teplotných Skupinách"
FIG_X_LABEL       =   "Priemer Častice [mm]"
FIG_Y_LABEL       =   "Rýchlosť Častice [m/s]"

FIG_GRID_STYLE    =   ":"

#======================================================
# ==                SCRIPT COMMANDS                  ==
#======================================================

# === Load CSV ===
df = pd.read_csv(CSV_FILENAME, sep=CSV_SEPARATOR)

# === Clean column names ===
df.columns = [col.strip() for col in df.columns]

# === Drop missing data & provide filtering ===
df = df.dropna(subset=[COL_DIAMETER, COL_VELOCITY, COL_TEMPERATURE])
df = df[df[COL_VELOCITY] >= LIM_VELOCITY_MIN]
df = df[(df[COL_DIAMETER] >= LIM_DIAMETER_MIN) & (df[COL_DIAMETER] <= LIM_DIAMETER_MAX)]

# === Group temperatures by 2°C bins ===
df[COL_TEMP_GROUPED] = (df[COL_TEMPERATURE] // TEMP_BIN_SIZE * TEMP_BIN_SIZE).astype(int)

# Find most frequent temperature group
most_common_temp = df[COL_TEMP_GROUPED].value_counts().idxmax()
count = df[COL_TEMP_GROUPED].value_counts().max()
print(f"The most frequent temperature group is {most_common_temp}–{most_common_temp+1}°C with {count} samples.")

# === Prepare plot ===
plt.figure(figsize=(FIG_WIDTH, FIG_HEIGHT))
unique_temps = sorted(df[COL_TEMP_GROUPED].unique())
colors = plt.cm.plasma(np.linspace(0, 1, len(unique_temps)))

# === Scatter + power fit for each temperature group ===
trendline_handle = None
for color, temp in zip(colors, unique_temps):
    group = df[df[COL_TEMP_GROUPED] == temp]
    plt.scatter(group[COL_DIAMETER], group[COL_VELOCITY],
                color=color, label=f"{temp}-{temp+1} °C", alpha=FIG_DATA_ALPHA, s=FIG_DATA_SIZE)

    # Power-law fit only for the most frequent group
    if temp == most_common_temp and len(group) >= 2:
        x = group[COL_DIAMETER].values
        y = group[COL_VELOCITY].values
        coeffs = np.polyfit(np.log(x), np.log(y), 1)
        b = coeffs[0]
        a = np.exp(coeffs[1])
        xfit = np.linspace(x.min(), x.max(), FIT_SAMPLE_COUNT)
        yfit = a * xfit ** b
        trendline_handle, = plt.plot(xfit, yfit, color=FIG_TR_COLOR, linewidth=FIG_TR_SIZE,
                                     label=f"Mocninová Aproximácia {temp}-{temp+1} °C")
        print(f"Power fit @ {temp}-{temp+1}°C: v = {a:.4e} * d^{b:.4f}")

plt.title(FIG_TITLE)
plt.xlabel(FIG_X_LABEL)
plt.ylabel(FIG_Y_LABEL)
plt.legend(title=FIG_LEGEND_TITLE, fontsize=FIG_LEGEND_SIZE, ncol=FIG_LEGEND_COLS)
plt.grid(True, linestyle=FIG_GRID_STYLE)
plt.tight_layout()
plt.show()