#!/usr/bin/env python3
"""
IZV cast1 projektu 2025

Autor: Roman Nečas
Login: xnecasr00

Detailni zadani projektu je v samostatnem projektu e-learningu.
Nezapomente na to, ze python soubory maji dane formatovani.

Muzete pouzit libovolnou vestavenou knihovnu a knihovny predstavene
na prednasce
"""

from typing import Any, Dict, List

import numpy as np
import requests
from bs4 import BeautifulSoup
from matplotlib.collections import LineCollection
from numpy.typing import NDArray

import matplotlib.pyplot as plt


def wave_inference_bad(
    x: NDArray[Any], y: NDArray[Any], sources: NDArray[Any], wavelength: float
) -> NDArray[Any]:
    """
    Referencni implementace, ktera je pomala a nevyuziva numpy efektivne;
    nezasahujte do ni!
    """
    k = 2 * np.pi / wavelength

    Z = np.zeros(x.shape + y.shape)
    for sx, sy in sources:
        for i in range(x.shape[0]):
            for j in range(y.shape[0]):
                R = np.sqrt((x[i] - sx) ** 2 + (y[j] - sy) ** 2)
                Z[j, i] += np.cos(k * R) / (1 + R)
    return Z


def wave_inference(
    x: NDArray[Any], y: NDArray[Any], sources: NDArray[Any], wavelength: float
) -> NDArray[Any]:
    """
    Efektivny vypocet interferencie vln pomocou NumPy vektorizacie.

    Vypocita amplitudu vlny Z pre 2D suradnicove polia x a y z viacerych
    zdrojov vln. Pouziva optimalizovany pristup s predalokovanym polom
    a iteraciou cez zdroje, co je rychlejsie pre maly pocet zdrojov.

    Args:
        x: 1D pole x suradnic
        y: 1D pole y suradnic
        sources: 2D pole tvaru (n, 2) s poziciami zdrojov [sx, sy]
        wavelength: vlnova dlzka vlny

    Returns:
        2D pole tvaru (len(y), len(x)) obsahujuce amplitudy vln
    """
    k = 2 * np.pi / wavelength

    # Vytvorenie meshgrid pre x a y suradnice
    X, Y = np.meshgrid(x, y)

    # Predalokacia vysledneho pola (rychlejsie a menej pamate)
    Z = np.zeros_like(X)

    # Iteracia cez zdroje - pre maly pocet zdrojov rychlejsie
    # nez vytvarat 3D tensor (nizsie rezijne naklady)
    for sx, sy in sources:
        # Vypocet vzdialenosti od zdroja
        # np.hypot je rychlejsi a presnejsi nez sqrt(dx**2 + dy**2)
        R = np.hypot(X - sx, Y - sy)
        # Pripocitanie prispevku amplitudy z tohto zdroja
        Z += np.cos(k * R) / (1 + R)

    return Z


def plot_wave(
    Z: NDArray[Any],
    x: NDArray[Any],
    y: NDArray[Any],
    show_figure: bool = False,
    save_path: str | None = None,
):
    """
    Vizualizacia interferencie vln ako heatmapa.

    Vytvori farebnu mapu amplitud vln so spravnymi popismi osi,
    nadpisom a color barom.

    Args:
        Z: 2D pole amplitud vln
        x: 1D pole x suradnic
        y: 1D pole y suradnic
        show_figure: ak True, zobrazi graf pomocou plt.show()
        save_path: ak je zadane, ulozi graf do tejto cesty

    Returns:
        None
    """
    plt.figure(figsize=(8, 6))

    # Vytvorenie heatmapy s rozsahom podla suradnic
    extent = [x.min(), x.max(), y.min(), y.max()]
    im = plt.imshow(
        Z,
        extent=extent,
        origin='lower',
        cmap='viridis',
        aspect='auto'
    )

    # Pridanie colorbaru s popisom
    cbar = plt.colorbar(im)
    cbar.set_label('Amplituda vlny')

    # Nastavenie popisov a nadpisu
    plt.xlabel('X pozice')
    plt.ylabel('Y pozice')
    plt.title('Vlnové pole')

    # Ulozenie ak je zadana cesta
    if save_path is not None:
        plt.savefig(save_path)

    # Zobrazenie ak je pozadovane
    if show_figure:
        plt.show()

    plt.close()


def generate_sinus(show_figure: bool = False, save_path: str | None = None):
    """
    Generovanie vizualizacie sinusovych a kosinusovych funkcii.

    Vytvori graf s dvoma podgrafmi:
    - Horny: sin(x) a cos(x) s vyplnenou plochou medzi nimi
    - Dolny: minimum (ciarkovane) a maximum (farebne podla zdroja)
      oboch funkcii

    Args:
        show_figure: ak True, zobrazi graf pomocou plt.show()
        save_path: ak je zadane, ulozi graf do tejto cesty

    Returns:
        None
    """
    # Generovanie x hodnot od 0 do 4pi
    x = np.linspace(0, 4 * np.pi, 1000)
    sin_vals = np.sin(x)
    cos_vals = np.cos(x)

    # Vytvorenie grafu s 2 podgrafmi so zdielanymi osami x a y
    _, (ax1, ax2) = plt.subplots(
        2, 1, figsize=(10, 8), sharex=True, sharey=True
    )

    # Prvy podgraf: sin a cos s vyplnenou plochou medzi nimi
    ax1.plot(x, sin_vals, label='sin(x)', color='blue')
    ax1.plot(x, cos_vals, label='cos(x)', color='orange')
    ax1.fill_between(x, sin_vals, cos_vals, alpha=0.3, color='green')
    ax1.set_ylabel('f(x)')
    ax1.grid(True, alpha=0.3)
    ax1.set_ylim(-1.5, 1.5)

    # Druhy podgraf: minimum (ciarkovane) a maximum (farebne)
    min_vals = np.minimum(sin_vals, cos_vals)
    max_vals = np.maximum(sin_vals, cos_vals)

    # Vykreslenie minima ciarkovane
    ax2.plot(x, min_vals, '--', color='gray', label='min')

    # Urcenie, ktora funkcia je maximalna v kazdom bode
    # True kde je cos max, False kde je sin max
    cos_is_max = cos_vals >= sin_vals

    # Vykreslenie maxima s farbami podla zdroja pomocou LineCollection
    # Oranzova kde je cos max, modra kde je sin max
    # Vytvorenie segmentov: body (x, y) spajane do useciek
    points = np.array([x, max_vals]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)

    # Farby pre kazdy segment podla zdroja maxima
    colors = np.where(cos_is_max[:-1], 'orange', 'blue')

    # Vytvorenie a pridanie LineCollection
    lc = LineCollection(segments, colors=colors, linewidth=2)
    ax2.add_collection(lc)

    ax2.set_ylabel('f(x)')
    ax2.grid(True, alpha=0.3)

    # Nastavenie znaciek na x-ovej osi s LaTeX pi notaciou
    pi_ticks = np.array([0, np.pi/2, np.pi, 3*np.pi/2, 2*np.pi,
                         5*np.pi/2, 3*np.pi, 7*np.pi/2, 4*np.pi])
    pi_labels = ['0', r'$\frac{\pi}{2}$', r'$\pi$', r'$\frac{3\pi}{2}$',
                 r'$2\pi$', r'$\frac{5\pi}{2}$', r'$3\pi$',
                 r'$\frac{7\pi}{2}$', r'$4\pi$']
    ax2.set_xticks(pi_ticks)
    ax2.set_xticklabels(pi_labels)
    ax2.set_xlabel('x')

    plt.tight_layout()

    # Ulozenie ak je zadana cesta
    if save_path is not None:
        plt.savefig(save_path)

    # Zobrazenie ak je pozadovane
    if show_figure:
        plt.show()

    plt.close()


def download_data() -> Dict[str, List[Any]]:
    """
    Stiahnutie dat meteorologickych stanic z web stranky.

    Stiahne a parsuje HTML tabulku z
    https://ehw.fit.vutbr.cz/izv/st_zemepis_cz.html
    obsahujucu informacie o ceskych meteorologickych staniciach.

    Returns:
        Slovnik s klucmi 'positions', 'lats', 'longs', 'heights'
        obsahujuci zoznamy nazvov stanic (str), zemepisnych sirok (float),
        zemepisnych dlzok (float) a vysok (float).
    """
    def _clean_numeric(text: str, remove_degree: bool = True) -> str:
        """Pomocna funkcia na cistenie numerickych retazcov."""
        # Vytvorenie prekladovej tabulky pre rychle nahradenie znakov
        trans_dict = {',': '.', '\xa0': '', ' ': ''}
        if remove_degree:
            trans_dict['°'] = ''
        trans_table = str.maketrans(trans_dict)
        return text.strip().translate(trans_table)

    # Skutocne data su v st_zemepis_cz.html
    url = "https://ehw.fit.vutbr.cz/izv/st_zemepis_cz.html"

    # Stiahnutie stranky
    response = requests.get(url)
    response.encoding = 'utf-8'

    # Parsovanie HTML
    soup = BeautifulSoup(response.text, 'html.parser')

    # Inicializacia vysledneho slovnika
    data = {
        'positions': [],
        'lats': [],
        'longs': [],
        'heights': []
    }

    # Najdenie vsetkych tabuliek a ziskanie datovej tabulky
    tables = soup.find_all('table')
    if len(tables) < 2:
        return data

    # Data su v druhej tabulke (prva je hlavicka)
    data_table = tables[1]

    # Najdenie vsetkych riadkov v tabulke
    rows = data_table.find_all('tr')

    # Spracovanie kazdeho riadku (preskocit hlavickovy riadok)
    for row in rows[1:]:
        cells = row.find_all('td')

        # Potrebujeme aspon 7 buniek podla struktury tabulky:
        # [nazov stanice, ikona linku, lat stupne, lat dms,
        #  long stupne, long dms, vyska]
        if len(cells) >= 7:
            # Extrahovanie nazvu stanice z prvej bunky (obsahuje <strong> tag)
            station_name_tag = cells[0].find('strong')
            if station_name_tag:
                position = station_name_tag.get_text(strip=True)
            else:
                position = cells[0].get_text(strip=True)

            # Cistenie numerickych hodnot pomocou optimalizovanej funkcie
            lat_str = _clean_numeric(cells[2].get_text(strip=True))
            long_str = _clean_numeric(cells[4].get_text(strip=True))
            height_str = _clean_numeric(cells[6].get_text(strip=True))

            # Konverzia na prislusne typy
            try:
                lat = float(lat_str)
                long = float(long_str)
                height = float(height_str)

                # Pridanie do zoznamov
                data['positions'].append(position)
                data['lats'].append(lat)
                data['longs'].append(long)
                data['heights'].append(height)
            except (ValueError, IndexError):
                # Preskocenie riadkov s neplatnymi datami
                continue

    return data


if __name__ == "__main__":
    # Priklad pouzitia - demonstracia implementovanych funkcii
    X = np.linspace(-10, 10, 200)
    Y = np.linspace(-10, 10, 200)
    A = wave_inference(X, Y, np.array([[-3, 0], [3, 0], [0, 4]]), 2)
    plot_wave(A, X, Y, show_figure=False, save_path="wave_example.png")
    generate_sinus(show_figure=False, save_path="sinus_example.png")
    print("Vygenerovane prikladove grafy: wave_example.png, sinus_example.png")