import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots


def stacked_plot(data):
    """Create a stacked plot with mean power on top and a 2D heatmap below.

    Parameters
    ----------
    data : array-like
        Input data array. Will be squeezed to remove singleton dimensions.
        Rows are interpreted as samples and columns as range/frequency bins.

    Returns
    -------
    plotly.graph_objects.Figure
        A two-row figure: the top panel shows the mean absolute power across
        columns, and the bottom panel shows a heatmap of the absolute values.
    """
    data = np.squeeze(data)
    mean_dp = np.mean(np.abs(data), axis=1)

    fig = make_subplots(rows=2, cols=1, row_heights=[0.3, 0.7], shared_xaxes=True,
                        vertical_spacing=0.01)

    fig.add_trace(go.Scatter(y=mean_dp, name='Mean Power'), row=1, col=1)
    fig.add_trace(go.Heatmap(z=np.abs(data).T, showscale=False, name='Heat Map'), row=2, col=1)

    fig.update_layout(title='Mean DP Power and 2D Map', autosize=True)
    fig.update_xaxes(visible=False, row=2, col=1)
    fig.update_yaxes(visible=False, row=2, col=1)

    return fig


def noise_mean(data):
    """Estimate the noise floor as the trimmed mean of absolute values.

    Sorts the flattened absolute data and discards the bottom 10% and top 10%
    before computing the mean, making the estimate robust to outliers and
    strong targets.

    Parameters
    ----------
    data : array-like
        Input data array of any shape.

    Returns
    -------
    float
        Mean of the central 80% of the sorted absolute values.
    """
    sorted_data = np.sort(np.abs(data.flatten()))
    cutoff_up_index = int(len(sorted_data) * 0.9)
    cutoff_down_index = int(len(sorted_data) * 0.1)
    trimmed_data = sorted_data[cutoff_down_index:cutoff_up_index]

    return np.mean(trimmed_data)


def calculate_cdf(data):
    """Compute the empirical cumulative distribution function (CDF) of the data.

    Parameters
    ----------
    data : array-like
        Input data array of any shape. Will be flattened before processing.

    Returns
    -------
    tuple[numpy.ndarray, numpy.ndarray]
        A ``(sorted_data, cdf)`` tuple where ``sorted_data`` contains the
        sorted values and ``cdf`` contains the corresponding CDF probabilities
        in the range (0, 1].
    """
    sorted_data = np.sort(data.flatten())
    cdf = np.arange(1, len(sorted_data) + 1) / len(sorted_data)
    return (sorted_data, cdf)


def plot_cdfs(data_list, labels):
    """Plot the empirical CDFs of multiple datasets on a single figure.

    Parameters
    ----------
    data_list : list of array-like
        List of data arrays to plot. Each array can have any shape and will
        be flattened internally by :func:`calculate_cdf`.
    labels : list of str
        Legend labels corresponding to each entry in ``data_list``.

    Returns
    -------
    plotly.graph_objects.Figure
        A figure with one CDF line per dataset.
    """
    fig = go.Figure()
    for data, label in zip(data_list, labels):
        sorted_data, cdf = calculate_cdf(data)
        fig.add_trace(go.Scatter(x=sorted_data, y=cdf, mode='lines', name=label))
    fig.update_layout(title='CDF of Data', xaxis_title='Value', yaxis_title='CDF', autosize=True)
    return fig