"""
The code below is adapted from the partitura library, and is therefore 
licenced under their Apache-2.0 licence. https://github.com/CPJKU/partitura. 
We make no alteration to the functionality of their code, but adapt it 
to work with our internal Melody object. All credit for this code remains 
with the original authors.

This module contains methods for estimation pitch spelling using the ps13 algorithm.

References
----------
Meredith, D. (2006). "The ps13 Pitch Spelling Algorithm". Journal
of New Music Research, 35(2):121.
"""
import numpy as np
from collections import namedtuple
from typing import List
from .representations import Melody

# from partitura.musicanalysis.utils import prepare_notearray


__all__ = ["estimate_spelling"]

ChromamorpheticPitch = namedtuple(
    "ChromamorpheticPitch", "chromatic_pitch morphetic_pitch"
)

STEPS = np.array(["A", "B", "C", "D", "E", "F", "G"])
UND_CHROMA = np.array([0, 2, 3, 5, 7, 8, 10], dtype=int)
ALTER = np.array(["n", "#", "b"])


def ensure_notearray(note_info):
    """Return a structured note array with required fields.

    Accepts either a numpy structured array or a Melody instance (converted via melody_to_note_array).

    Raises ValueError/TypeError if the input is not supported or missing
    required fields.
    """
    # Convert Melody directly
    if isinstance(note_info, Melody):
        return melody_to_note_array(note_info)

    # Pass through structured numpy arrays that have the needed fields
    if isinstance(note_info, np.ndarray) and note_info.dtype.names is not None:
        field_names = set(note_info.dtype.names)
        if "pitch" not in field_names:
            raise ValueError("note_info must contain a 'pitch' field")
        has_beat = {"onset_beat", "duration_beat"}.issubset(field_names)
        has_sec = {"onset_sec", "duration_sec"}.issubset(field_names)
        if not (has_beat or has_sec):
            raise ValueError(
                "note_info must contain either ('onset_beat','duration_beat') or ('onset_sec','duration_sec')"
            )
        return note_info

    raise TypeError(
        "Unsupported note_info type. Provide a structured numpy array with required fields or a Melody instance."
    )


def get_time_units_from_note_array(note_array):
    """Determine the onset and duration field names to use.

    Preference order: beat units first, then seconds. If neither pair is
    present, raises a ValueError.
    """
    if not (isinstance(note_array, np.ndarray) and note_array.dtype.names is not None):
        raise TypeError("note_array must be a structured numpy array with named fields")

    names = set(note_array.dtype.names)
    if {"onset_beat", "duration_beat"}.issubset(names):
        return "onset_beat", "duration_beat"
    if {"onset_sec", "duration_sec"}.issubset(names):
        return "onset_sec", "duration_sec"
    if {"onset", "duration"}.issubset(names):
        return "onset", "duration"
    raise ValueError(
        "Could not determine onset/duration fields. Expected beat or seconds fields to be present."
    )


def estimate_spelling(note_info, method="ps13s1", **kwargs):
    """Estimate pitch spelling using the ps13 algorithm [4]_, [5]_.

    Parameters
    ----------
    note_info : structured array, `Part` or `PerformedPart`
        Note information as a `Part` or `PerformedPart` instances or
        as a structured array. If it is a structured array, it has to
        contain the fields generated by the `note_array` properties
        of `Part` or `PerformedPart` objects. If the array contains
        onset and duration information of both score and performance,
        (e.g., containing both `onset_beat` and `onset_sec`), the score
        information will be preferred.
    method : {'ps13s1'}
         Pitch spelling algorithm. More methods will be added.
    **kwargs
        Keyword arguments for the algorithm specified in `method`.

    Returns
    -------
    spelling : structured array
        Array with pitch spellings. The fields are 'step', 'alter' and
        'octave'

    References
    ----------
    .. [4] Meredith, D. (2006). "The ps13 Pitch Spelling Algorithm". Journal
           of New Music Research, 35(2):121.
    .. [5] Meredith, D. (2019). "RecurSIA-RRT: Recursive translatable
           point-set pattern discovery with removal of redundant translators".
           12th International Workshop on Machine Learning and Music. Würzburg,
           Germany.

    """
    if method == "ps13s1":
        ps = ps13s1
    else:
        raise ValueError(f"Unsupported pitch spelling method: {method}")
    step, alter, octave = ps(ensure_notearray(note_info), **kwargs)

    spelling = np.empty(
        len(step), dtype=[("step", "U1"), ("alter", int), ("octave", int)]
    )

    spelling["step"] = step
    spelling["alter"] = alter
    spelling["octave"] = octave

    return spelling


def _prepare_tempo_map(tempo_changes):
    """Return sorted tempo change arrays and cumulative beats at each change.

    Ensures the first change starts at time 0.0.
    """
    if not tempo_changes:
        tempo_times = np.array([0.0], dtype=float)
        tempo_bpms = np.array([120.0], dtype=float)
    else:
        # Sort tempo changes by time and ensure (0.0, tempo0) exists
        tempo_changes = sorted(tempo_changes, key=lambda x: float(x[0]))
        if tempo_changes[0][0] > 0.0:
            tempo_changes = [(0.0, float(tempo_changes[0][1]))] + tempo_changes
        tempo_times = np.array([float(t) for t, _ in tempo_changes], dtype=float)
        tempo_bpms = np.array([float(bpm) for _, bpm in tempo_changes], dtype=float)

    # Cumulative beats at the start of each segment
    # For segments [t_i, t_{i+1}), beats added are (t_{i+1}-t_i) * bpm_i / 60
    cum_beats = np.zeros_like(tempo_times, dtype=float)
    if len(tempo_times) > 1:
        dt = np.diff(tempo_times)
        seg_beats = dt * tempo_bpms[:-1] / 60.0
        cum_beats[1:] = np.cumsum(seg_beats)

    return tempo_times, tempo_bpms, cum_beats


def _seconds_to_beats(times_sec, tempo_changes):
    """Convert array-like seconds to beats using piecewise-constant tempo map.

    times_sec can be scalar or array-like. Returns a NumPy array of beats.
    """
    tempo_times, tempo_bpms, cum_beats = _prepare_tempo_map(tempo_changes)

    times = np.asarray(times_sec, dtype=float)
    # Find segment index for each time: last tempo change with time <= t
    idx = np.searchsorted(tempo_times, times, side="right") - 1
    idx = np.clip(idx, 0, len(tempo_times) - 1)

    beats_from_segment_start = (times - tempo_times[idx]) * (tempo_bpms[idx] / 60.0)
    beats = cum_beats[idx] + beats_from_segment_start
    return beats


def melody_to_note_array(melody: "Melody"):
    """Convert a Melody instance into a Partitura-compatible note array.

    The returned structured array contains the fields required by the
    pitch spelling pipeline used here: 'pitch', 'onset_sec', 'duration_sec'.

    Parameters
    ----------
    melody : Melody
        Our internal Melody representation with pitches, starts and ends in seconds.

    Returns
    -------
    np.ndarray
        Structured array with dtype [('pitch', int), ('onset_sec', float), ('duration_sec', float)].
    """
    pitches: List[int] = melody.pitches
    starts: List[float] = melody.starts
    ends: List[float] = melody.ends

    n = min(len(pitches), len(starts), len(ends))
    if n == 0:
        return np.empty(0, dtype=[("pitch", int), ("onset_sec", float), ("duration_sec", float)])

    durations = [max(0.0, (ends[i] - starts[i])) for i in range(n)]

    note_array = np.empty(
        n,
        dtype=[
            ("pitch", int),
            ("onset_sec", float),
            ("duration_sec", float),
            ("onset_beat", float),
            ("duration_beat", float),
        ],
    )
    note_array["pitch"] = pitches[:n]
    note_array["onset_sec"] = starts[:n]
    note_array["duration_sec"] = durations

    # Compute beat-based timing using the melody's tempo map (supports tempo changes)
    starts_beats = _seconds_to_beats(note_array["onset_sec"], melody.tempo_changes)
    ends_beats = _seconds_to_beats(np.array(starts[:n]) + np.array(durations), melody.tempo_changes)
    note_array["onset_beat"] = starts_beats
    note_array["duration_beat"] = np.maximum(0.0, ends_beats - starts_beats)

    return note_array


def estimate_spelling_from_melody(melody: "Melody", method: str = "ps13s1", **kwargs):
    """Estimate pitch spelling directly from a Melody instance.

    Parameters
    ----------
    melody : Melody
        Our internal Melody object.
    method : {'ps13s1'}
        Pitch spelling algorithm.

    Returns
    -------
    np.ndarray
        Structured array with fields 'step', 'alter', 'octave'. Same format as estimate_spelling.
    """
    note_array = melody_to_note_array(melody)
    return estimate_spelling(note_array, method=method, **kwargs)


def format_spelled_notes(spelling_array) -> List[str]:
    """Format a spelling array (step, alter, octave) into note name strings.

    Returns e.g. ['C#4', 'Bb3', 'F4'] matching the order of the input notes.
    """
    names: List[str] = []
    for step, alter, octave in spelling_array:
        if alter > 0:
            acc = "#"
        elif alter < 0:
            acc = "b"
        else:
            acc = ""
        names.append(f"{step}{acc}{int(octave)}")
    return names


def ps13s1(note_array, K_pre=10, K_post=40):
    """
    ps13s1 Pitch Spelling Algorithm
    """

    onset_unit, _ = get_time_units_from_note_array(note_array)
    pitch_sort_idx = note_array["pitch"].argsort()

    onset_sort_idx = np.argsort(
        note_array[pitch_sort_idx][onset_unit], kind="mergesort"
    )

    sort_idx = pitch_sort_idx[onset_sort_idx]

    re_idx = sort_idx.argsort()  # o_idx[sort_idx]

    sorted_ocp = np.column_stack(
        (
            note_array[sort_idx][onset_unit],
            chromatic_pitch_from_midi(note_array[sort_idx]["pitch"]),
        )
    )

    # n = len(sorted_ocp)
    # ChromaList
    chroma_array = compute_chroma_array(sorted_ocp=sorted_ocp)
    # ChromaVectorList
    chroma_vector_array = compute_chroma_vector_array(
        chroma_array=chroma_array, K_pre=K_pre, K_post=K_post
    )
    morph_array = compute_morph_array(
        chroma_array=chroma_array, chroma_vector_array=chroma_vector_array
    )

    morphetic_pitch = compute_morphetic_pitch(sorted_ocp, morph_array)

    step, alter, octave = p2pn(
        sorted_ocp[:, 1],
        morphetic_pitch.reshape(
            -1,
        ),
    )
    # sort back pitch names
    step = step[re_idx]
    alter = alter[re_idx]
    octave = octave[re_idx]

    return step, alter, octave


def chromatic_pitch_from_midi(midi_pitch):
    return midi_pitch - 21


def chroma_from_chromatic_pitch(chromatic_pitch):
    return np.mod(chromatic_pitch, 12)


def pitch_class_from_chroma(chroma):
    return np.mod(chroma - 3, 12)


def compute_chroma_array(sorted_ocp):
    return chroma_from_chromatic_pitch(sorted_ocp[:, 1]).astype(int)


def compute_chroma_vector_array(chroma_array, K_pre, K_post):
    """
    Computes the chroma frequency distribution within the context surrounding
    each note.
    """
    n = len(chroma_array)
    chroma_vector = np.zeros(12, dtype=int)

    for i in range(np.minimum(n, K_post)):
        chroma_vector[chroma_array[i]] = 1 + chroma_vector[chroma_array[i]]

    chroma_vector_list = [chroma_vector.copy()]

    for i in range(1, n):
        if i + K_post <= n:
            chroma_vector[chroma_array[i + K_post - 1]] = (
                1 + chroma_vector[chroma_array[i + K_post - 1]]
            )

        if i - K_pre > 0:
            chroma_vector[chroma_array[i - K_pre - 1]] = (
                chroma_vector[chroma_array[i - K_pre - 1]] - 1
            )

        chroma_vector_list.append(chroma_vector.copy())

    return np.array(chroma_vector_list)


def compute_morph_array(chroma_array, chroma_vector_array):
    n = len(chroma_array)
    # Line 1: Initialize morph array
    morph_array = np.empty(n, dtype=int)

    # Compute m0
    # Line 2
    init_morph = np.array([0, 1, 1, 2, 2, 3, 4, 4, 5, 5, 6, 6], dtype=int)
    # Line 3
    c0 = chroma_array[0]
    # Line 4
    m0 = init_morph[c0]

    # Line 5
    morph_int = np.array([0, 1, 1, 2, 2, 3, 3, 4, 5, 5, 6, 6], dtype=int)

    # Lines 6-8
    tonic_morph_for_tonic_chroma = np.mod(
        m0 - morph_int[np.mod(c0 - np.arange(12), 12)], 7
    )

    # Line 10
    tonic_chroma_set_for_morph = [[] for i in range(7)]

    # Line 11
    morph_strength = np.zeros(7, dtype=int)

    # Line 12
    for j in range(n):
        # Lines 13-15 (skipped line 9, since we do not need to
        # initialize morph_for_tonic_chroma)
        morph_for_tonic_chroma = np.mod(
            morph_int[np.mod(chroma_array[j] - np.arange(12), 12)]
            + tonic_morph_for_tonic_chroma,
            7,
        )
        # Lines 16-17
        tonic_chroma_set_for_morph = [[] for i in range(7)]

        # Line 18
        for m in range(7):
            # Line 19
            for ct in range(12):
                # Line 20
                if morph_for_tonic_chroma[ct] == m:
                    # Line 21
                    tonic_chroma_set_for_morph[m].append(ct)

        # Line 22
        for m in range(7):
            # Line 23
            morph_strength[m] = sum(
                [chroma_vector_array[j, ct] for ct in tonic_chroma_set_for_morph[m]]
            )

        # Line 24
        morph_array[j] = np.argmax(morph_strength)

    return morph_array


def compute_ocm_chord_list(sorted_ocp, chroma_array, morph_array):
    # Lines 1-3
    ocm_array = np.column_stack((sorted_ocp[:, 0], chroma_array, morph_array)).astype(
        int
    )

    # Alternative implementation of lines 4--9
    unique_onsets = np.unique(ocm_array[:, 0])
    unique_onset_idxs = [np.where(ocm_array[:, 0] == u) for u in unique_onsets]
    ocm_chord_list = [ocm_array[uix] for uix in unique_onset_idxs]

    return ocm_chord_list


def compute_morphetic_pitch(sorted_ocp, morph_array):
    """
    Compute morphetic pitch

    Parameters
    ----------
    sorted_ocp : array
       Sorted array of (onset in beats, chromatic pitch)
    morph_array : array
       Array of morphs

    Returns
    -------
    morphetic_pitch : array
        Morphetic pitch of the notes
    """
    n = len(sorted_ocp)
    chromatic_pitch = sorted_ocp[:, 1]
    morph = morph_array.reshape(-1, 1)

    morph_oct_1 = np.floor(chromatic_pitch / 12.0).astype(int)

    morph_octs = np.column_stack((morph_oct_1, morph_oct_1 + 1, morph_oct_1 - 1))

    chroma = np.mod(chromatic_pitch, 12)

    mps = morph_octs + (morph / 7)

    cp = (morph_oct_1 + (chroma / 12)).reshape(-1, 1)

    diffs = abs(cp - mps)

    best_morph_oct = morph_octs[np.arange(n), diffs.argmin(1)]

    morphetic_pitch = (
        morph.reshape(
            -1,
        )
        + 7 * best_morph_oct
    )

    return morphetic_pitch


def p2pn(c_pitch, m_pitch):
    """
    Chromamorphetic pitch to pitch name

    Parameters
    ----------
    c_pitch : int or array
        Chromatic pitch.
    m_pitch : int or array
        Morphetic pitch.

    Returns
    -------
    step : str or array
        Note name (step)
    alter : int or array
        Alteration(s) of the notes. 1 is sharp, -1 is flat and 0 is natural
    octave : int or array
        Octave
    """
    morph = np.mod(m_pitch, 7)

    step = STEPS[morph]
    undisplaced_chroma = UND_CHROMA[morph]

    # displacement in paper
    alter = c_pitch - 12 * np.floor(m_pitch / 7.0) - undisplaced_chroma

    asa_octave = np.floor(m_pitch / 7)

    if isinstance(morph, (int, float)):
        if morph > 1:
            asa_octave += 1
    else:
        asa_octave[morph > 1] += 1

    return step, alter, asa_octave