import sys
import os
from datetime import datetime
import numpy as np
import plotly.graph_objects as go

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import config

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def visualize_fdm(T_matrix, x_vals, t_vals):
    """Produce three HTML visualisations for the FDM solver output.

    Parameters
    ----------
    T_matrix : np.ndarray, shape (NX, NT)
        Temperature field: T_matrix[i, n] = T(x_i, t_n).
    x_vals : np.ndarray, shape (NX,)
        Spatial node positions.
    t_vals : np.ndarray, shape (NT,)
        Time values.
    """
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    out_dir = os.path.join(BASE_DIR, 'results', 'fdm', timestamp)
    os.makedirs(out_dir, exist_ok=True)

    # ------------------------------------------------------------------
    # 1. heatmap.html: plot statico T(x,t) + striscia 1D animata
    #    Due figure Plotly indipendenti nello stesso file HTML per
    #    evitare conflitti tra animazioni e assi multipli.
    # ------------------------------------------------------------------
    zmin_val = float(np.min(T_matrix))
    zmax_val = float(np.max(T_matrix))

    # Subsample frames (shared with animation below)
    n_frames = len(t_vals)
    max_frames = 200
    step = max(1, n_frames // max_frames)
    frame_indices = list(range(0, n_frames, step))

    # --- Figura A: heatmap 2D statica (identica all'originale) ---
    z_data = T_matrix.T  # (NT, NX)
    z_center = float(np.mean(T_matrix))
    z_abs = float(np.max(np.abs(T_matrix - z_center))) or 1.0
    fig_static = go.Figure(go.Heatmap(
        z=z_data,
        x=x_vals,
        y=t_vals,
        colorscale='RdBu_r',
        zmin=z_center - z_abs,
        zmax=z_center + z_abs,
        colorbar=dict(title='T [°C]'),
    ))
    fig_static.update_layout(
        title='FDM — Heat Equation  T(x,t)',
        xaxis_title='x [m]',
        yaxis_title='t [s]',
        height=480,
        margin=dict(t=50, b=50, l=70, r=90),
    )

    # --- Figura B: striscia 1D animata ---
    strip_frames = [
        go.Frame(
            data=[go.Heatmap(
                z=T_matrix[:, idx].reshape(1, -1),
                x=x_vals,
                y=[0],
                colorscale='Jet',
                zmin=zmin_val,
                zmax=zmax_val,
                showscale=True,
                colorbar=dict(title='T [°C]', thickness=18),
            )],
            name=str(idx),
            layout=go.Layout(title_text=f't = {t_vals[idx]:.3f} s'),
        )
        for idx in frame_indices
    ]

    fig_strip = go.Figure(
        data=[go.Heatmap(
            z=T_matrix[:, frame_indices[0]].reshape(1, -1),
            x=x_vals,
            y=[0],
            colorscale='Jet',
            zmin=zmin_val,
            zmax=zmax_val,
            showscale=True,
            colorbar=dict(title='T [°C]', thickness=18),
        )],
        frames=strip_frames,
        layout=go.Layout(
            title=f't = {t_vals[frame_indices[0]]:.3f} s',
            xaxis=dict(title='x [m]'),
            yaxis=dict(showticklabels=False, showgrid=False,
                       zeroline=False, fixedrange=True),
            height=300,
            margin=dict(t=80, b=130, l=70, r=110),
            updatemenus=[dict(
                type='buttons',
                showactive=False,
                y=1.22, x=0.5, xanchor='center',
                buttons=[
                    dict(label='▶ Play', method='animate',
                         args=[None, dict(frame=dict(duration=40, redraw=True),
                                          fromcurrent=True, mode='immediate')]),
                    dict(label='⏸ Pause', method='animate',
                         args=[[None], dict(frame=dict(duration=0, redraw=False),
                                            mode='immediate')]),
                ],
            )],
            sliders=[dict(
                steps=[
                    dict(method='animate',
                         args=[[str(idx)], dict(mode='immediate',
                                                frame=dict(duration=0, redraw=True))],
                         label=f'{t_vals[idx]:.2f}')
                    for idx in frame_indices
                ],
                transition=dict(duration=0),
                x=0.05, y=0, len=0.9,
                currentvalue=dict(prefix='t = ', font=dict(size=14)),
            )],
        ),
    )

    # Scrivi entrambe le figure in un unico file HTML
    html_static = fig_static.to_html(full_html=False, include_plotlyjs='cdn')
    html_strip = fig_strip.to_html(full_html=False, include_plotlyjs=False)
    heatmap_path = os.path.join(out_dir, 'heatmap.html')
    with open(heatmap_path, 'w', encoding='utf-8') as _f:
        _f.write('<!DOCTYPE html><html><head><meta charset="utf-8"></head><body>\n')
        _f.write(html_static)
        _f.write('\n')
        _f.write(html_strip)
        _f.write('\n</body></html>')
    print(f"Heatmap saved        → {heatmap_path}")

    # ------------------------------------------------------------------
    # 2. Animated profile T(x) evolving in time
    # ------------------------------------------------------------------
    L = config.L

    y_min = float(np.min(T_matrix)) - 1.0
    y_max = float(np.max(T_matrix)) + 1.0

    vline_shapes = [
        dict(type='line', x0=0, x1=0, y0=y_min, y1=y_max,
             line=dict(color='grey', dash='dash', width=1)),
        dict(type='line', x0=L, x1=L, y0=y_min, y1=y_max,
             line=dict(color='grey', dash='dash', width=1)),
    ]

    frames = []
    for idx in frame_indices:
        frames.append(go.Frame(
            data=[go.Scatter(
                x=x_vals,
                y=T_matrix[:, idx],
                mode='lines',
                line=dict(color='royalblue', width=2),
                name='FDM',
            )],
            name=str(idx),
            layout=go.Layout(
                title_text=f'FDM  |  t = {t_vals[idx]:.3f} s',
            ),
        ))

    fig_anim = go.Figure(
        data=[go.Scatter(
            x=x_vals,
            y=T_matrix[:, frame_indices[0]],
            mode='lines',
            line=dict(color='royalblue', width=2),
            name='FDM',
        )],
        layout=go.Layout(
            title=f'FDM  |  t = {t_vals[frame_indices[0]]:.3f} s',
            xaxis=dict(title='x [m]', range=[-0.02 * L, 1.02 * L]),
            yaxis=dict(title='T [°C]', range=[y_min, y_max]),
            shapes=vline_shapes,
            updatemenus=[dict(
                type='buttons',
                showactive=False,
                y=1.15,
                x=0.5,
                xanchor='center',
                buttons=[
                    dict(
                        label='▶ Play',
                        method='animate',
                        args=[None, dict(
                            frame=dict(duration=40, redraw=False),
                            fromcurrent=True,
                            mode='immediate',
                        )],
                    ),
                    dict(
                        label='⏸ Pause',
                        method='animate',
                        args=[[None], dict(
                            frame=dict(duration=0, redraw=False),
                            mode='immediate',
                        )],
                    ),
                ],
            )],
            sliders=[dict(
                steps=[
                    dict(
                        method='animate',
                        args=[[str(idx)], dict(
                            mode='immediate',
                            frame=dict(duration=0, redraw=False),
                        )],
                        label=f'{t_vals[idx]:.2f}',
                    )
                    for idx in frame_indices
                ],
                transition=dict(duration=0),
                x=0.05,
                y=0,
                len=0.9,
                currentvalue=dict(prefix='t = ', font=dict(size=14)),
            )],
        ),
        frames=frames,
    )

    anim_path = os.path.join(out_dir, 'animation.html')
    fig_anim.write_html(anim_path)
    print(f"Animation saved      → {anim_path}")

    # ------------------------------------------------------------------
    # 3. Time evolution at fixed spatial points
    # ------------------------------------------------------------------
    fixed_fractions = [0.0, 0.25, 0.5, 0.75, 1.0]
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd']

    fig_ts = go.Figure()

    for frac, color in zip(fixed_fractions, colors):
        x_target = frac * L
        idx_x = int(np.argmin(np.abs(x_vals - x_target)))
        label = f'x = {x_vals[idx_x]:.2f} m'
        fig_ts.add_trace(go.Scatter(
            x=t_vals,
            y=T_matrix[idx_x, :],
            mode='lines',
            name=label,
            line=dict(color=color, width=2),
        ))

    # "Heat ON" vertical dashed line at t = T_STEP
    t_step = config.T_STEP
    t_min_val = float(np.min(t_vals))
    t_max_val = float(np.max(t_vals))
    T_min_val = float(np.min(T_matrix)) - 1.0
    T_max_val = float(np.max(T_matrix)) + 1.0

    fig_ts.add_shape(
        type='line',
        x0=t_step, x1=t_step,
        y0=T_min_val, y1=T_max_val,
        line=dict(color='red', dash='dash', width=1.5),
    )
    fig_ts.add_annotation(
        x=t_step,
        y=T_max_val,
        text='Heat ON',
        showarrow=False,
        yanchor='top',
        font=dict(color='red', size=11),
    )

    fig_ts.update_layout(
        title='FDM — Temperature evolution at fixed points',
        xaxis=dict(title='t [s]', range=[t_min_val, t_max_val]),
        yaxis=dict(title='T [°C]', range=[T_min_val, T_max_val]),
        legend=dict(x=0.01, y=0.99),
        height=480,
    )

    ts_path = os.path.join(out_dir, 'time_series.html')
    fig_ts.write_html(ts_path)
    print(f"Time-series saved    → {ts_path}")