plot_results.py

#!/usr/bin/env python3
"""
plot_results.py — Interaktywny kreator wykresów z plików wyników symulacji.

Użycie:
    python plot_results.py

Przepływ:
  1. Wybierz plik wyników z katalogu results/
  2. Wybierz kolumnę osi X  (domyślnie: x_m → wyświetlana w mm)
  3. Wybierz kolumny osi Y  (jedna lub kilka, np. "7,8,9" lub "7-9")
  4. Wykres wyświetla się w oknie
  5. Po zamknięciu okna: opcja zapisu → opcja kolejnego wykresu
"""

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

RESULTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results')

# ---------------------------------------------------------------------------
# Kolumny [m] do automatycznej konwersji → mm na wykresie
# (tylko jednostkowe długości, nie powierzchnie _m2 ani prędkości _m_s)
# ---------------------------------------------------------------------------
_MM_COLUMNS = {'x_m', 'r_m'}

# ---------------------------------------------------------------------------
# Czytelne etykiety osi dla każdej kolumny
# ---------------------------------------------------------------------------
_LABELS = {
    'x_m':              'x [mm]',
    'r_m':              'r [mm]',
    'A_m2':             'A [m²]',
    'gamma':            'Wykładnik adiabaty γ [-]',
    'Cpcg_J_kgK':      'Cp gazu [J/(kg·K)]',
    'Prcg_gas':         'Pr gazu [-]',
    'molar_mass_kg_mol':'Masa molarna gazu [kg/mol]',
    'Rs_J_kgK':         'Stała gazowa właściwa Rs [J/(kg·K)]',
    'M':                'Liczba Macha M [-]',
    'N_M2':             'N = M² [-]',
    'P_Pa':             'Ciśnienie gazu p [Pa]',
    'T_K':              'Temperatura gazu T [K]',
    'T_aw_K':           'T adiabatyczna ścianki T_aw [K]',
    'eta_Pa_s':         'Lepkość dynamiczna η [Pa·s]',
}

# ---------------------------------------------------------------------------
# Grupy kolumn do czytelnego wyświetlania w menu
# ---------------------------------------------------------------------------
_GROUPS = [
    ('Geometria',          ['x_m', 'r_m', 'A_m2']),
    ('Właściwości gazu',   ['gamma', 'Cpcg_J_kgK', 'Prcg_gas',
                            'molar_mass_kg_mol', 'Rs_J_kgK']),
    ('Przepływ gazowy',    ['M', 'N_M2', 'P_Pa', 'T_K', 'T_aw_K']),
    ('Inne',               ['eta_Pa_s']),
]


# ===========================================================================
# Pomocnicze
# ===========================================================================

def _label(col):
    """Zwróć czytelną etykietę osi dla kolumny."""
    return _LABELS.get(col, col.replace('_', ' '))


def _get_values(df, col):
    """Zwróć wartości kolumny z konwersją m → mm tam gdzie potrzeba."""
    vals = df[col].values.astype(float)
    if col in _MM_COLUMNS:
        vals = vals * 1000.0
    return vals


# ===========================================================================
# Wczytywanie pliku
# ===========================================================================

def _read_metadata(filepath):
    """Parsuj blok metadanych ('# klucz: wartość') z pliku CSV."""
    meta = {}
    with open(filepath, 'r', encoding='utf-8') as fh:
        for line in fh:
            if not line.startswith('#'):
                break
            line = line[1:].strip()
            if ':' in line:
                key, _, val = line.partition(':')
                meta[key.strip()] = val.strip()
    return meta


def _list_result_files():
    """Zwróć posortowaną listę plików CSV z katalogu results/."""
    if not os.path.isdir(RESULTS_DIR):
        return []
    return sorted(
        os.path.join(RESULTS_DIR, f)
        for f in os.listdir(RESULTS_DIR)
        if f.endswith('.csv')
    )


# ===========================================================================
# Interaktywne wybory
# ===========================================================================

def select_result_file():
    """Wyświetl listę plików i pozwól użytkownikowi wybrać jeden."""
    files = _list_result_files()

    if not files:
        print(f"\n  Brak plików wyników w: {RESULTS_DIR}")
        print("  Uruchom najpierw symulację (python main.py).")
        return None

    print("\n" + "=" * 65)
    print("  WYBÓR PLIKU WYNIKÓW")
    print("=" * 65)

    for i, fp in enumerate(files, 1):
        meta = _read_metadata(fp)
        fname  = os.path.basename(fp)
        date   = meta.get('date', '?')
        conv   = meta.get('converged', '?')
        iters  = meta.get('iterations', '?')
        params = meta.get('params_name', '?')
        print(f"  [{i:2d}] {fname}")
        print(f"        parametry: {params}  |  data: {date}")
        print(f"        zbieżność: {conv}  |  iteracje: {iters}")

    print("-" * 65)

    while True:
        try:
            raw = input(f"Wybierz plik (1–{len(files)}): ").strip()
            idx = int(raw) - 1
            if 0 <= idx < len(files):
                return files[idx]
            print(f"  Podaj liczbę od 1 do {len(files)}.")
        except ValueError:
            print("  Podaj numer pliku.")
        except (EOFError, KeyboardInterrupt):
            return None


def _build_index_map(df):
    """
    Przypisz numery do kolumn zgodnie z grupami.
    Zwraca dict {numer: nazwa_kolumny}.
    """
    cols_available = set(df.columns)
    index_map = {}
    idx = 1

    for _, group_cols in _GROUPS:
        for col in group_cols:
            if col in cols_available:
                index_map[idx] = col
                idx += 1

    # Kolumny spoza znanych grup (przyszłościowe rozszerzenia)
    assigned = {c for _, gc in _GROUPS for c in gc}
    for col in df.columns:
        if col not in assigned:
            index_map[idx] = col
            idx += 1

    return index_map


def display_columns(df, index_map):
    """Wyświetl dostępne kolumny w pogrupowanej tabeli."""
    cols_available = set(df.columns)
    reverse = {v: k for k, v in index_map.items()}

    print()
    print("  Dostępne parametry:")
    print("  " + "-" * 61)

    for group_name, group_cols in _GROUPS:
        present = [c for c in group_cols if c in cols_available]
        if not present:
            continue
        print(f"  {group_name}:")
        for col in present:
            num = reverse[col]
            note = '  (→ mm)' if col in _MM_COLUMNS else ''
            print(f"    [{num:2d}] {col:22s}  {_label(col)}{note}")

    # Nieprzypisane
    assigned = {c for _, gc in _GROUPS for c in gc}
    extra = [c for c in df.columns if c not in assigned]
    if extra:
        print("  Inne:")
        for col in extra:
            num = reverse[col]
            print(f"    [{num:2d}] {col:22s}  {_label(col)}")

    print("  " + "-" * 61)


def _parse_selection(raw, index_map, multi=True):
    """
    Parsuj ciąg numerów: "7,8,9" lub "7-9" lub "7, 9-12, 15".
    Zwraca listę nazw kolumn w kolejności wyboru.
    """
    selected = []
    for token in raw.replace(' ', '').split(','):
        if '-' in token and multi:
            parts = token.split('-', 1)
            try:
                a, b = int(parts[0]), int(parts[1])
                for n in range(a, b + 1):
                    if n in index_map and index_map[n] not in selected:
                        selected.append(index_map[n])
            except ValueError:
                pass
        else:
            try:
                n = int(token)
                if n in index_map and index_map[n] not in selected:
                    selected.append(index_map[n])
            except ValueError:
                pass
    return selected


def select_x_axis(df, index_map):
    """
    Wybór kolumny osi X.
    Domyślnie: x_m (wyświetlana w mm).
    """
    default_col = 'x_m' if 'x_m' in df.columns else list(index_map.values())[0]
    default_num = next((k for k, v in index_map.items() if v == default_col), 1)

    print(f"\n  Oś X — wybierz numer  (domyślnie [{default_num}] = {_label(default_col)}):")
    try:
        raw = input(f"  X [{default_num}]: ").strip()
    except (EOFError, KeyboardInterrupt):
        return default_col

    if not raw:
        return default_col

    result = _parse_selection(raw, index_map, multi=False)
    if result:
        return result[0]

    print(f"  Nieprawidłowy wybór — użyto domyślnego: {default_col}")
    return default_col


def select_y_axes(df, index_map):
    """
    Wybór kolumn osi Y (jedna lub więcej).
    Wpisz numery oddzielone przecinkami lub zakresy, np. "7,8,9" / "7-9".
    """
    print("  Oś Y — wybierz numery kolumn (przecinki lub zakresy, np. 7,8 lub 7-9):")
    while True:
        try:
            raw = input("  Y: ").strip()
        except (EOFError, KeyboardInterrupt):
            return []

        if not raw:
            print("  Podaj co najmniej jeden numer kolumny.")
            continue

        result = _parse_selection(raw, index_map, multi=True)
        if result:
            return result

        print("  Nie rozpoznano żadnej kolumny. Spróbuj ponownie.")


# ===========================================================================
# Tworzenie wykresu
# ===========================================================================

def make_chart(df, x_col, y_cols, meta):
    """
    Wygeneruj wykres matplotlib i zwróć obiekt fig.

    Jeśli wybrano tylko jedną serię Y, etykieta osi Y jest pełna.
    Jeśli wybrano wiele serii, etykieta osi Y to 'Wartość' i używana jest legenda.
    """
    x_vals  = _get_values(df, x_col)
    x_label = _label(x_col)

    fig, ax = plt.subplots(figsize=(11, 6))

    for col in y_cols:
        y_vals = _get_values(df, col)
        ax.plot(x_vals, y_vals, linewidth=1.8, label=_label(col))

    ax.set_xlabel(x_label, fontsize=12)

    if len(y_cols) == 1:
        ax.set_ylabel(_label(y_cols[0]), fontsize=12)
    else:
        ax.set_ylabel('Wartość', fontsize=12)
        ax.legend(fontsize=10, loc='best')

    # Tytuł z metadanych pliku
    params_name = meta.get('params_name', '')
    date        = meta.get('date', '?')
    conv        = meta.get('converged', '?')
    iters       = meta.get('iterations', '?')

    title    = f"OpenEngine — {params_name}"
    subtitle = f"data analizy: {date}  |  zbieżność: {conv}  |  iteracje: {iters}"
    ax.set_title(f"{title}\n{subtitle}", fontsize=11)

    ax.grid(True, alpha=0.35)
    fig.tight_layout()
    return fig


# ===========================================================================
# Zapis wykresu
# ===========================================================================

def ask_save_chart(fig, params_name):
    """Zapytaj użytkownika czy zapisać wykres i w jakim formacie."""
    try:
        ans = input("\n  Zapisać wykres do pliku? (t/n, domyślnie: n): ").strip().lower()
    except (EOFError, KeyboardInterrupt):
        return

    if ans != 't':
        return

    default_name = f"{params_name}_plot"
    try:
        name = input(f"  Nazwa pliku bez rozszerzenia (domyślnie '{default_name}'): ").strip()
    except (EOFError, KeyboardInterrupt):
        name = ''

    if not name:
        name = default_name

    safe_name = "".join(c for c in name if c.isalnum() or c in ('_', '-', '.'))
    if not safe_name:
        safe_name = default_name

    print("  Format zapisu:  [1] PNG   [2] PDF   [3] SVG")
    try:
        fmt = input("  Format [1]: ").strip()
    except (EOFError, KeyboardInterrupt):
        fmt = '1'

    ext = {' ': '.png', '': '.png', '1': '.png', '2': '.pdf', '3': '.svg'}.get(fmt, '.png')

    os.makedirs(RESULTS_DIR, exist_ok=True)
    out_path = os.path.join(RESULTS_DIR, safe_name + ext)
    fig.savefig(out_path, dpi=150, bbox_inches='tight')
    print(f"  Zapisano: {out_path}")


# ===========================================================================
# Główna pętla programu
# ===========================================================================

def main():
    print("\n" + "=" * 65)
    print("  OPENENGINE — KREATOR WYKRESÓW")
    print("=" * 65)

    # --- Wybór pliku ---
    filepath = select_result_file()
    if filepath is None:
        print("\n  Anulowano.")
        return

    df, meta = pd.read_csv(filepath, comment='#'), _read_metadata(filepath)
    # (ponowne wczytanie przez pandas — upewnia się że typy są numeryczne)
    df = pd.read_csv(filepath, comment='#')

    index_map = _build_index_map(df)

    fname = os.path.basename(filepath)
    print(f"\n  Plik:       {fname}")
    print(f"  Parametry:  {meta.get('params_name', '?')}")
    print(f"  Data:       {meta.get('date', '?')}")
    print(f"  Siatka:     {len(df)} punktów  |  Kolumny: {len(df.columns)}")

    params_name = meta.get('params_name', 'results')

    # --- Pętla wykresów ---
    while True:
        display_columns(df, index_map)

        x_col  = select_x_axis(df, index_map)
        y_cols = select_y_axes(df, index_map)

        if not y_cols:
            print("  Brak wybranych kolumn Y — pomijam.")
        else:
            print(f"\n  Rysuję: X = {_label(x_col)}")
            print(f"          Y = {', '.join(_label(c) for c in y_cols)}")

            fig = make_chart(df, x_col, y_cols, meta)
            plt.show(block=True)          # czeka na zamknięcie okna

            ask_save_chart(fig, params_name)
            plt.close(fig)

        # --- Kolejny wykres? ---
        print()
        try:
            again = input("  Zrobić kolejny wykres? (t/n, domyślnie: n): ").strip().lower()
        except (EOFError, KeyboardInterrupt):
            again = 'n'

        if again != 't':
            print("\n  Do widzenia!\n")
            break


if __name__ == '__main__':
    main()