"""Full-dimensional analogy arithmetic and shared linear projections."""
import numpy as np
import plotly.graph_objects as go
from plotly.colors import qualitative


def cosine(a, b):
    a, b = np.asarray(a, dtype=float), np.asarray(b, dtype=float)
    denominator = np.linalg.norm(a) * np.linalg.norm(b)
    return float(np.clip(np.dot(a, b) / denominator, -1, 1)) if denominator > 0 and np.isfinite(denominator) else None


def analyse_pairs(pairs, target, expected=None):
    """Each pair is (A, B); apply their mean offset to C, optionally compare D."""
    if not pairs:
        raise ValueError('Add at least one example pair.')
    arrays = [np.asarray(v, dtype=float) for pair in pairs for v in pair] + [np.asarray(target, dtype=float)]
    if expected is not None:
        arrays.append(np.asarray(expected, dtype=float))
    if any(v.ndim != 1 or v.shape != arrays[0].shape or not np.isfinite(v).all() for v in arrays):
        raise ValueError('All inputs must be finite vectors with matching dimensions.')
    offsets = np.array([np.asarray(b, dtype=float) - np.asarray(a, dtype=float) for a, b in pairs])
    offset = offsets.mean(axis=0)
    prediction = np.asarray(target, dtype=float) + offset
    agreement = [cosine(v, offset) for v in offsets]
    result = {'offset': offset, 'prediction': prediction, 'offset_length': float(np.linalg.norm(offset)),
              'example_agreement': float(np.mean([x for x in agreement if x is not None])) if any(x is not None for x in agreement) else None}
    if expected is not None:
        observed = np.asarray(expected, dtype=float) - np.asarray(target, dtype=float)
        alignment = cosine(offset, observed)
        result.update(direction_cosine=alignment,
                      angle_degrees=float(np.degrees(np.arccos(alignment))) if alignment is not None else None,
                      observed_length=float(np.linalg.norm(observed)),
                      length_ratio=float(np.linalg.norm(observed) / np.linalg.norm(offset)) if np.linalg.norm(offset) > 0 else None,
                      prediction_error=float(np.linalg.norm(prediction - expected)))
    return result


def fit_projection(points, dimensions):
    points = np.asarray(points, dtype=float)
    if points.ndim != 2 or not len(points) or not np.isfinite(points).all() or dimensions not in (2, 3):
        raise ValueError('Projection needs finite points and two or three components.')
    mean = points.mean(axis=0)
    centered = points - mean
    _, singular, vt = np.linalg.svd(centered, full_matrices=False)
    basis = np.zeros((dimensions, points.shape[1]))
    tolerance = max(centered.shape) * np.finfo(float).eps * (singular[0] if len(singular) else 0.)
    rank = int(np.count_nonzero(singular > tolerance))
    count = min(dimensions, rank)
    basis[:count] = vt[:count]
    total = float(np.dot(singular, singular))
    retained = float(np.dot(singular[:count], singular[:count]) / total) if total > 0 else 1.
    return mean, basis, retained


def relationship_figure(analysis, dimensions, extra_vectors=None, extra_labels=None):
    """Fit one PCA to the example endpoints; project predictions/candidates with it."""
    points, point_labels, point_codes, point_keys = [], [], [], []
    for index, pair in enumerate(analysis['pairs'], 1):
        for letter in ('a', 'b'):
            points.append(pair[letter]['values'])
            point_labels.append(pair[letter]['label'])
            point_codes.append(letter.upper() + (str(index) if len(analysis['pairs']) > 1 else ''))
            point_keys.append(pair[letter]['key'])
    points.append(analysis['target']['values'])
    point_labels.append(analysis['target']['label'])
    point_codes.append('C')
    point_keys.append(analysis['target']['key'])
    if analysis['expected'] is not None:
        points.append(analysis['expected']['values'])
        point_labels.append(analysis['expected']['label'])
        point_codes.append('D')
        point_keys.append(analysis['expected']['key'])
    mean, basis, retained = fit_projection(points, dimensions)
    project = lambda values: (np.asarray(values) - mean) @ basis.T
    coords = project(points)
    predicted = project(analysis['metrics']['prediction'])
    c_index = 2 * len(analysis['pairs'])
    palette = qualitative.Plotly
    fig = go.Figure()
    scatter = go.Scatter if dimensions == 2 else go.Scatter3d

    def add_points(coordinates, **kwargs):
        axes = {'x': coordinates[:, 0], 'y': coordinates[:, 1]}
        if dimensions == 3:
            axes['z'] = coordinates[:, 2]
        fig.add_trace(scatter(**axes, **kwargs))

    extent = max(float(np.linalg.norm(np.ptp(np.vstack([coords, predicted]), axis=0))), .01)

    def arrow(start, end, label, color, dash='solid', showlegend=True, group=None):
        group = group or label
        add_points(np.array([start, end]), mode='lines', name=label, showlegend=showlegend,
                   legendgroup=group,
                   line={'color': color, 'width': 4 if dimensions == 3 else 2, 'dash': dash},
                   hoverinfo='name')
        if np.linalg.norm(end - start) < 1e-12:
            return
        if dimensions == 2:
            fig.add_annotation(x=end[0], y=end[1], ax=start[0], ay=start[1], xref='x', yref='y',
                               axref='x', ayref='y', text='', showarrow=True, arrowhead=3,
                               arrowwidth=2, arrowcolor=color)
        else:
            delta = end - start
            fig.add_trace(go.Cone(x=[end[0]], y=[end[1]], z=[end[2]], u=[delta[0]], v=[delta[1]], w=[delta[2]],
                                 anchor='tip', sizemode='absolute', sizeref=extent * .045,
                                 colorscale=[[0, color], [1, color]], showscale=False, hoverinfo='skip',
                                 legendgroup=group, showlegend=False))

    for i in range(len(analysis['pairs'])):
        arrow(coords[i * 2], coords[i * 2 + 1], 'Example offset' if len(analysis['pairs']) == 1 else 'Example offsets',
              palette[0], showlegend=i == 0, group='examples')
    arrow(coords[c_index], predicted, 'C + mean(B − A)', palette[2], 'dash')
    if analysis['expected'] is not None:
        arrow(coords[c_index], coords[c_index + 1], 'C → D', palette[1])
        arrow(predicted, coords[c_index + 1], 'Prediction error', palette[3], 'dot')
    add_points(coords, mode='markers+text', text=point_codes, textposition='top center',
               customdata=point_labels, name='Input vectors',
               marker={'size': 7 if dimensions == 2 else 4, 'color': palette[0]},
               hovertemplate='%{text}: %{customdata}<extra></extra>')
    add_points(predicted[None, :], mode='markers+text', text=['Predicted'], textposition='bottom center',
               name='Predicted D', marker={'size': 10 if dimensions == 2 else 6, 'symbol': 'diamond', 'color': palette[2]},
               hovertemplate='Predicted D = C + mean(B − A)<extra></extra>')
    extra = [key for key in (extra_vectors or {}) if key not in point_keys]
    if extra:
        extra_coords = project([extra_vectors[key] for key in extra])
        add_points(extra_coords, mode='markers', customdata=[extra_labels[key] for key in extra],
                   name='Other chart vectors', marker={'size': 5 if dimensions == 2 else 3, 'color': palette[4]},
                   hovertemplate='%{customdata}<extra></extra>')
    fig.update_layout(template='plotly_dark', autosize=True,
                      title={'text': f'PCA · {retained:.1%} of input variance retained', 'font': {'size': 13}, 'x': .01, 'y': .98, 'yanchor': 'top'},
                      margin={'l': 20, 'r': 20, 't': 72, 'b': 25},
                      legend={'orientation': 'h', 'y': 1.02, 'yanchor': 'bottom'},
                      uirevision=f'relationship:{analysis["id"]}:{dimensions}')
    if dimensions == 2:
        fig.update_xaxes(title='PC 1')
        fig.update_yaxes(title='PC 2', scaleanchor='x', scaleratio=1)
    else:
        fig.update_layout(scene={'xaxis_title': 'PC 1', 'yaxis_title': 'PC 2', 'zaxis_title': 'PC 3', 'aspectmode': 'data'})
    return fig