import os
from typing import Any
import logging

import matplotlib.pyplot as plt
import numpy as np
import polars as pl

import latdraw
from . import convert
from .lattice import Beamline
from latdraw import lattice

logging.basicConfig(level=logging.INFO)
LOG = logging.getLogger(__name__)


def _get_lattice(lattice: Any) -> lattice.Beamline:
    if isinstance(lattice, Beamline):
        return lattice
    return convert.coerce_beamline(lattice)


def subplots_with_lattice(
    lattice: str | lattice.Beamline | pl.DataFrame | None,
    s_offset=0,
    nrows: int = 1,
    gridspec_kw=None,
    **kwargs,
):
    # Plot goes at the top
    pattern = [lattice]
    pattern.extend(nrows * [None])
    return subplots_with_lattices(pattern, s_offset=s_offset, **kwargs)


def subplots_with_lattices(
    pattern, s_offset=0, **draw_kwargs
) -> tuple[plt.Figure, list[plt.Axes]]:
    pattern = np.array(pattern, dtype=object)

    height_ratios = np.full_like(pattern, 1.0, dtype=float)
    # Get indices of where machines should be plotted
    indices = [index for (index, value) in enumerate(pattern) if value is not None]
    height_ratios[indices] = 0.25

    the_gridspec_kw = {"height_ratios": height_ratios, "hspace": 0.05}

    fig, axes = plt.subplots(
        nrows=len(pattern), sharex=True, gridspec_kw=the_gridspec_kw
    )

    for lattice, ax in zip(pattern, axes):
        if lattice is None:
            continue

        lattice = _get_lattice(lattice)
        latdraw.draw(fig, ax, lattice, s_offset=s_offset, **draw_kwargs)

        ax.set_yticks([], [])

        ax.tick_params(
            top=False,
            bottom=False,
            left=False,
            right=False,
            labelleft=False,
            labelbottom=False,
        )

        ax.spines["left"].set_visible(False)
        ax.spines["right"].set_visible(False)

        # ax.spines['top'].set_visible(True)
        # ax.spines['bottom'].set_visible(False)

        ax.set_ylim(-0.25, 0.25)

    return fig, axes


def plot_optics(
    some_beamline, some_optics, title=""
) -> tuple[plt.Figure, tuple[plt.Axes, plt.Axes, plt.Axes]]:
    bl = convert.coerce_beamline(some_beamline)
    optics = latdraw.optics.coerce_optics(some_optics)

    fig, axes = subplots_with_lattice(bl, nrows=2)

    fig, (mx, ax1, ax2) = fig, axes

    ax1.plot(optics["s"], optics["beta_x"], label=r"$\beta_x$")
    ax1.plot(optics["s"], optics["beta_y"], label=r"$\beta_y$")
    ax1.legend()
    beta_label(ax1)

    ax2.plot(optics["s"], optics["dy"], label=r"$D_x$")
    ax2.plot(optics["s"], optics["dy"], label=r"$D_y$")
    dispersion_label(ax2)
    s_label(ax2)
    ax2.legend()

    mx.set_title(title)

    return fig, axes


def simple_figure(
    some_beamline, title="", **drawkwargs
) -> tuple[plt.Figure, tuple[plt.Axes, plt.Axes]]:
    bl = convert.coerce_beamline(some_beamline)

    fig, axes = subplots_with_lattice(bl, nrows=1, **drawkwargs)
    s_label(axes[-1])
    axes[0].set_title(title)
    # fig.suptitle(title)
    return fig, axes


def two_axes_figure(
    some_beamline, title=""
) -> tuple[plt.Figure, tuple[plt.Axes, plt.Axes, plt.Axes]]:
    bl = convert.coerce_beamline(some_beamline)

    fig, axes = subplots_with_lattice(bl, nrows=2)
    s_label(axes[-1])
    axes[0].set_title(title)
    # fig.suptitle(title)
    return fig, axes


def three_axes_figure(
    some_beamline, title=""
) -> tuple[plt.Figure, tuple[plt.Axes, plt.Axes, plt.Axes, plt.Axes]]:
    bl = convert.coerce_beamline(some_beamline)

    fig, axes = subplots_with_lattice(bl, nrows=3)
    s_label(axes[-1])
    axes[0].set_title(title)
    return fig, axes


def four_axes_figure(
    some_beamline: Beamline, title=""
) -> tuple[plt.Figure, tuple[plt.Axes, plt.Axes, plt.Axes, plt.Axes]]:
    bl = convert.coerce_beamline(some_beamline)

    fig, axes = subplots_with_lattice(bl, nrows=4)
    s_label(axes[-1])
    axes[0].set_title(title)
    return fig, axes


def beta_label(ax: plt.Axes, subscript: str = "") -> None:
    if not subscript:
        ax.set_ylabel(r"$\beta$ / m")
    else:
        ax.set_ylabel(rf"$\beta_{subscript}$ / m")


def alpha_label(ax: plt.Axes, subscript: str = "") -> None:
    if not subscript:
        ax.set_ylabel(r"$\alpha$ / m")
    else:
        ax.set_ylabel(rf"$\alpha_{subscript}$")


def dispersion_label(ax: plt.Axes, subscript: str = "") -> None:
    if not subscript:
        ax.set_ylabel(r"$D$ / m")
    else:
        ax.set_ylabel(rf"$D_{{{subscript}}}$ / m")


def s_label(ax: plt.Axes) -> None:
    ax.set_xlabel(r"$s$ / m")


def energy_label(ax: plt.Axes, unit: str = "GeV") -> None:
    ax.set_ylabel(f"$E$ / {unit}")
