from smallmuse.temporals import NoteList, TemporalNote

# Parameters
QUANTIZATION = 8  # positions per bar
TICKS_PER_BEAT = 2
STEPS_PER_BAR = QUANTIZATION
MAX_TIME_STEP = 32
BAR_LIMIT = .5  # ±1 bar displacement allowed
DISSONANT_INTERVALS = {6, 13}


def _load_cp_model():
    try:
        from ortools.sat.python import cp_model
    except ImportError as exc:
        raise ImportError(
            "OR-Tools is required for temporal note optimization. "
            "Install it with `python -m pip install ortools`."
        ) from exc
    return cp_model

# Note class
# class Note:
#     def __init__(self, pitch, start_time, duration):
#         self.pitch = pitch
#         self.start_time = start_time
#         self.duration = duration

# Generate notes
# notes = []
# for i in range(NUM_NOTES):
#     pitch = random.randint(48, 72)
#     start_time = i * 0.5  # uniform spacing (beats)
#     duration = 0.5
#     notes.append(Note(pitch, start_time, duration))

def optimize(note_list):
    cp_model = _load_cp_model()
    # Model
    notes = note_list.notes
    NUM_NOTES = len(notes)
    model = cp_model.CpModel()
    original_steps = [int(note.start_time * TICKS_PER_BEAT) for note in notes]

    starts = []
    for i in range(NUM_NOTES):
        lower = max(0, original_steps[i] - int(BAR_LIMIT * STEPS_PER_BAR))
        upper = min(MAX_TIME_STEP - 1, original_steps[i] + int(BAR_LIMIT * STEPS_PER_BAR))
        var = model.NewIntVar(lower, upper, f'start_{i}')
        starts.append(var)

    # Order constraint
    for i in range(NUM_NOTES - 1):
        model.Add(starts[i] <= starts[i + 1])

    # Dissonance constraint
    for i in range(NUM_NOTES):
        for j in range(i + 1, NUM_NOTES):
            interval = abs(notes[i].pitch - notes[j].pitch)
            if interval in DISSONANT_INTERVALS:
                model.Add(starts[i] != starts[j])

    # Objective: minimize squared pitch difference
    # interval_terms = [(notes[i + 1].pitch - notes[i].pitch) ** 2 for i in range(NUM_NOTES - 1)]
    # model.Minimize(sum(interval_terms))

    # Objective: minimize total displacement from original start times
    displacements = []
    for i in range(NUM_NOTES):
        diff = model.NewIntVar(0, MAX_TIME_STEP, f'diff_{i}')
        model.AddAbsEquality(diff, starts[i] - original_steps[i])
        displacements.append(diff)

    model.Minimize(sum(displacements))

    # Solve
    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    # Output
    result = NoteList()
    if status in [cp_model.OPTIMAL, cp_model.FEASIBLE]:
        for i in range(NUM_NOTES):
            result.add_note(TemporalNote(notes[i].pitch, int(solver.Value(starts[i]) / 2), notes[i].duration()))
            # print(f"Note {i}: pitch={notes[i].pitch}, "
            #       f"original={original_steps[i]}, "
            #       f"quantized={solver.Value(starts[i])}")
    else:
        print("No solution found.")
    result.join_notes()
    return result