#ifndef MARKOV_STRATEGY_H
#define MARKOV_STRATEGY_H

#include "MelodyStrategy.h"
#include <Arduino.h>

class MarkovStrategy : public MelodyStrategy {
public:
    void generate(Step (*sequence)[NUM_STEPS], int track, int numSteps, int* scaleNotes, int numScaleNotes, int seed) override {
        randomSeed(seed);
        if (numScaleNotes == 0) return;

        // Transition matrix: weights for moving from index i to index j
        // Max scale size is 12
        uint8_t matrix[12][12];
        
        // Initialize matrix with weighted probabilities
        for (int i = 0; i < numScaleNotes; i++) {
            for (int j = 0; j < numScaleNotes; j++) {
                // Base random weight
                matrix[i][j] = random(1, 8);
                
                // Musical heuristics: favor small intervals
                int dist = abs(i - j);
                if (dist == 0) matrix[i][j] += 8;  // Repeat note
                else if (dist == 1) matrix[i][j] += 12; // Stepwise motion
                else if (dist == 2) matrix[i][j] += 6;  // Thirds
                // Large jumps remain low probability
            }
        }

        int currentIdx = random(numScaleNotes);

        for (int i = 0; i < numSteps; i++) {
            // 20% chance of rest
            if (random(100) < 20) {
                sequence[track][i].note = -1;
                sequence[track][i].accent = false;
                sequence[track][i].tie = false;
                continue;
            }

            // Select next note based on matrix
            int totalWeight = 0;
            for (int j = 0; j < numScaleNotes; j++) {
                totalWeight += matrix[currentIdx][j];
            }

            int r = random(totalWeight);
            int nextIdx = 0;
            for (int j = 0; j < numScaleNotes; j++) {
                r -= matrix[currentIdx][j];
                if (r < 0) {
                    nextIdx = j;
                    break;
                }
            }
            currentIdx = nextIdx;

            // Determine Octave (weighted towards middle)
            int octave = 4;
            int rOct = random(100);
            if (rOct < 20) octave = 3;
            else if (rOct > 80) octave = 5;

            sequence[track][i].note = 12 * octave + scaleNotes[currentIdx];
            sequence[track][i].accent = (random(100) < 25);
            sequence[track][i].tie = (random(100) < 15);
        }
        randomSeed(micros());
    }

    void generateScale(int* scaleNotes, int& numScaleNotes) override {
        numScaleNotes = random(5, 8); // Pentatonic to Major
        for (int i = 0; i < 12; i++) {
            scaleNotes[i] = i; // Fill with all notes
        }
        // Shuffle
        for (int i = 0; i < 12; i++) {
            int j = random(12);
            int temp = scaleNotes[i];
            scaleNotes[i] = scaleNotes[j];
            scaleNotes[j] = temp;
        }
        sortArray(scaleNotes, numScaleNotes);
    }

    void mutate(Step (*sequence)[NUM_STEPS], int track, int numSteps, int* scaleNotes, int numScaleNotes) override {
        // Drift mutation: pick a note and move it stepwise in the scale
        int s = random(numSteps);
        if (sequence[track][s].note != -1) {
            int currentNoteVal = sequence[track][s].note % 12;
            int idx = -1;
            // Find index in scale
            for(int i=0; i<numScaleNotes; i++) {
                if (scaleNotes[i] == currentNoteVal) {
                    idx = i;
                    break;
                }
            }
            
            if (idx != -1) {
                // Move up or down 1 step in scale
                if (random(2) == 0) idx = (idx + 1) % numScaleNotes;
                else idx = (idx - 1 + numScaleNotes) % numScaleNotes;
                
                int octave = sequence[track][s].note / 12;
                sequence[track][s].note = 12 * octave + scaleNotes[idx];
            }
        } else {
            // Chance to fill a rest
            if (random(100) < 25) {
                 int octave = 3 + random(3);
                 sequence[track][s].note = 12 * octave + scaleNotes[random(numScaleNotes)];
                 sequence[track][s].accent = false;
                 sequence[track][s].tie = false;
            }
        }
    }

    const char* getName() override {
        return "Markov";
    }
};

#endif