/*
* BEGIN CLASS DECLARATIONS
*/

// 2D vector class to make some of the math easier
class Vec2D {
    constructor(x, y) {
        this.x = x;
        this.y = y;
    }

    magnitude() {
        return Math.sqrt(this.x * this.x + this.y * this.y);
    }

    angle() {
        return Math.atan2(this.y, this.x);
    }
}

// Swerve module class to represent a single wheel
class SwerveModule {
    constructor(x, y, name) {
        this.position = new Vec2D(x, y);
        this.velocity = new Vec2D(0, 0);
        this.angle = 0;
        this.speed = 0;
        this.name = name;
    }

    calculateState(velocityX, velocityY, turnSpeed, heading = 0) {
        // Take the requested speed and turn rate of the robot and calculate
        // speed and angle of this module to achieve it

        // Transform field-relative velocities to robot-relative velocities
        // by rotating the velocity vector by the negative of the robot's heading
        const cosHeading = Math.cos(-heading);
        const sinHeading = Math.sin(-heading);

        const robotVelX = velocityX * cosHeading - velocityY * sinHeading;
        const robotVelY = velocityX * sinHeading + velocityY * cosHeading;

        // Calculate rotation contribution (perpendicular to position vector)
        const rotX = -this.position.y * turnSpeed;
        const rotY = this.position.x * turnSpeed;

        // Combine translation and rotation (now in robot frame)
        this.velocity.x = robotVelX + rotX;
        this.velocity.y = robotVelY + rotY;

        // Calculate speed and angle (in robot frame)
        this.speed = this.velocity.magnitude();
        this.angle = this.velocity.angle();
    }
}

// Swerve drive class to represent the robot as a whole
class SwerveDrive {
    constructor(modulePositionsAndNames, robotName) {
        this.setModules(modulePositionsAndNames);
        this.setName(robotName);
        this.gyroHeading = 0; // Simulated gyro heading in radians
    }

    setName(robotName) {
        this.name = robotName;
    }

    setModules(modulePositionsAndNames) {
        // Take an array of module positions with a name and create an array of SwerveModule objects
        this.modules = modulePositionsAndNames.map(module =>
            new SwerveModule(module.x, module.y, module.name)
        );
    }

    updateHeading(turnSpeed, deltaTime = 0.01) {
        // Integrate turn speed to update gyro heading
        // turnSpeed is in radians/second, deltaTime is the time step
        this.gyroHeading += turnSpeed * deltaTime;

        // Normalize to -PI to PI range
        while (this.gyroHeading > Math.PI) this.gyroHeading -= 2 * Math.PI;
        while (this.gyroHeading < -Math.PI) this.gyroHeading += 2 * Math.PI;
    }

    drive(velocityX, velocityY, turnSpeed, maxModuleSpeed, deltaTime = 0.01) {
        // Update gyro heading first
        this.updateHeading(turnSpeed, deltaTime);

        // Take in a requested speeds and update every module
        this.modules.forEach(module =>
            module.calculateState(velocityX, velocityY, turnSpeed, this.gyroHeading)
        );

        // If any speeds exceed the max speed, normalize down so we don't effect movement direction
        const maxCalculated = Math.max(...this.modules.map(m => m.speed), 0);
        if (maxCalculated > maxModuleSpeed) {
            const scale = maxModuleSpeed / maxCalculated;
            this.modules.forEach(module => {
                module.velocity.x *= scale;
                module.velocity.y *= scale;
                module.speed = module.velocity.magnitude();
                module.angle = module.velocity.angle();
            });
        }
    }
}

// Preset robot generators
const PresetConfigs = {
    twoWheel: (size) => [
        { x: size / 2, y: 0, name: "Left" },
        { x: -size / 2, y: 0, name: "Right" }
    ],

    threeWheel: (size) => {
        const radius = size / 2;
        return [
            { x: radius * Math.cos(Math.PI / 2), y: radius * Math.sin(Math.PI / 2), name: "Front" },
            { x: radius * Math.cos(Math.PI / 2 + 2 * Math.PI / 3), y: radius * Math.sin(Math.PI / 2 + 2 * Math.PI / 3), name: "Back Left" },
            { x: radius * Math.cos(Math.PI / 2 + 4 * Math.PI / 3), y: radius * Math.sin(Math.PI / 2 + 4 * Math.PI / 3), name: "Back Right" }
        ];
    },

    fourWheelSquare: (size) => {
        const half = size / 2;
        return [
            { x: half, y: half, name: "FL" },
            { x: half, y: -half, name: "FR" },
            { x: -half, y: half, name: "BL" },
            { x: -half, y: -half, name: "BR" }
        ];
    },

    fourWheelRectangle: (size) => {
        const width = size * 0.5;
        const length = size;
        return [
            { x: length / 2, y: width / 2, name: "FL" },
            { x: length / 2, y: -width / 2, name: "FR" },
            { x: -length / 2, y: width / 2, name: "BL" },
            { x: -length / 2, y: -width / 2, name: "BR" }
        ];
    },

    sixWheel: (size) => {
        const radius = size / 2;
        const modules = [];
        for (let i = 0; i < 6; i++) {
            const angle = (Math.PI / 2) + (i * Math.PI / 3);
            modules.push({
                x: radius * Math.cos(angle),
                y: radius * Math.sin(angle),
                name: `Module ${i + 1}`
            });
        }
        return modules;
    },

    eightWheel: (size) => {
        const radius = size / 2;
        const modules = [];
        for (let i = 0; i < 8; i++) {
            const angle = (Math.PI / 2) + (i * Math.PI / 4);
            modules.push({
                x: radius * Math.cos(angle),
                y: radius * Math.sin(angle),
                name: `Module ${i + 1}`
            });
        }
        return modules;
    }
};

/*
* END CLASS DECLARATIONS
* BEGIN DOM VARIABLES
*/

// Get all control elements
const vxSlider = document.getElementById('vx-slider');
const vySlider = document.getElementById('vy-slider');
const omegaSlider = document.getElementById('omega-slider');
const maxSpeedSlider = document.getElementById('max-speed-slider');
const moduleCountInput = document.getElementById('module-count');

// Get all output elements
const vxOutput = document.getElementById('vx-value');
const vyOutput = document.getElementById('vy-value');
const omegaOutput = document.getElementById('omega-value');
const maxSpeedOutput = document.getElementById('max-speed-value');

// Get button elements
const resetBtn = document.getElementById('reset-btn');
const generateInputsBtn = document.getElementById('generate-inputs-btn');
const clearInputsBtn = document.getElementById('delete-inputs-btn');
const applyCustomBtn = document.getElementById('apply-custom-btn');

// Preset buttons
const preset2WheelBtn = document.getElementById('preset-2wheel');
const preset3WheelBtn = document.getElementById('preset-3wheel');
const preset4WheelBtn = document.getElementById('preset-4wheel');
const preset4RectBtn = document.getElementById('preset-4rect');
const preset6WheelBtn = document.getElementById('preset-6wheel');
const preset8WheelBtn = document.getElementById('preset-8wheel');

/*
* END DOM VARIABLES
* BEGIN LISTENER CODE
*/


vxSlider.addEventListener('input', (e) => {
    vxOutput.textContent = parseFloat(e.target.value);
});
vxOutput.textContent = parseFloat(vxSlider.value);

vySlider.addEventListener('input', (e) => {
    vyOutput.textContent = parseFloat(e.target.value);
});
vyOutput.textContent = parseFloat(vySlider.value);

omegaSlider.addEventListener('input', (e) => {
    omegaOutput.textContent = parseFloat(e.target.value);
});
omegaOutput.textContent = parseFloat(omegaSlider.value);

maxSpeedSlider.addEventListener('input', (e) => {
    maxSpeedOutput.textContent = parseFloat(e.target.value);
});
maxSpeedOutput.textContent = parseFloat(maxSpeedSlider.value);

resetBtn.addEventListener('click', (e) => {
    vxSlider.value = 0;
    vySlider.value = 0;
    omegaSlider.value = 0;

    vxOutput.textContent = parseFloat(vxSlider.value);
    vyOutput.textContent = parseFloat(vySlider.value);
    omegaOutput.textContent = parseFloat(omegaSlider.value);
});

// Preset button event listeners
preset2WheelBtn.addEventListener('click', () => {
    const positions = PresetConfigs.twoWheel(robotSize);
    robot.setModules(positions);
    robot.setName("2-Wheel Differential");
    createModuleDisplays(robot);
    updateModuleDisplays(robot);
});

preset3WheelBtn.addEventListener('click', () => {
    const positions = PresetConfigs.threeWheel(robotSize);
    robot.setModules(positions);
    robot.setName("3-Wheel Triangle");
    createModuleDisplays(robot);
    updateModuleDisplays(robot);
});

preset4WheelBtn.addEventListener('click', () => {
    const positions = PresetConfigs.fourWheelSquare(robotSize);
    robot.setModules(positions);
    robot.setName("4-Wheel Square");
    createModuleDisplays(robot);
    updateModuleDisplays(robot);
});

preset4RectBtn.addEventListener('click', () => {
    const positions = PresetConfigs.fourWheelRectangle(robotSize);
    robot.setModules(positions);
    robot.setName("4-Wheel Rectangle");
    createModuleDisplays(robot);
    updateModuleDisplays(robot);
});

preset6WheelBtn.addEventListener('click', () => {
    const positions = PresetConfigs.sixWheel(robotSize);
    robot.setModules(positions);
    robot.setName("6-Wheel Hexagon");
    createModuleDisplays(robot);
    updateModuleDisplays(robot);
});

preset8WheelBtn.addEventListener('click', () => {
    const positions = PresetConfigs.eightWheel(robotSize);
    robot.setModules(positions);
    robot.setName("8-Wheel Octogon");
    createModuleDisplays(robot);
    updateModuleDisplays(robot);
});

generateInputsBtn.addEventListener('click', () => {
    const count = parseInt(moduleCountInput.value);

    if (isNaN(count) || count < 2) {
        alert('Please enter a valid number of modules above or equal to 2.');
        return;
    }
    generateModuleInputs(count);
    applyCustomBtn.style.display = 'block';
});

clearInputsBtn.addEventListener('click', () => {
    generateModuleInputs(0);
    applyCustomBtn.style.display = 'none';
});

applyCustomBtn.addEventListener('click', () => {
    const container = document.getElementById('module-position-inputs');
    const moduleElements = container.childNodes;

    const customModules = [];
    for (let i = 0; i < moduleElements.length; i++) {
        const xInput = document.getElementById(`module-${i}-x`);
        const yInput = document.getElementById(`module-${i}-y`);
        const nameInput = document.getElementById(`module-${i}-name`);

        const x = parseFloat(xInput.value);
        const y = parseFloat(yInput.value);
        const name = nameInput.value.trim();

        customModules.push({ x, y, name });
    }

    robot.setModules(customModules);
    robot.setName("Custom Configuration");
    createModuleDisplays(robot);
    updateModuleDisplays(robot);
});

/*
* END LISTENER CODE
* BEGIN DYNAMIC DOM FUNCTIONS
*/

function generateModuleInputs(count) {
    const container = document.getElementById('module-position-inputs');
    container.innerHTML = ''; // Clear existing inputs

    for (let i = 0; i < count; i++) {
        const moduleFieldset = document.createElement('fieldset');
        moduleFieldset.className = 'module-input-group';
        moduleFieldset.innerHTML = `
        <legend>Module ${i + 1}</legend>

        <div class="control-group">
            <label for="module-${i}-name">Module Name</label>
            <input type="text" id="module-${i}-name" value="Module ${i + 1}" required>
        </div>
        <div class="control-group">
            <label for="module-${i}-x">X Position (pixels)</label>
            <input type="number" id="module-${i}-x" step="1" value="0" required>
        </div>
        <div class="control-group">
            <label for="module-${i}-y">Y Position (pixels)</label>
            <input type="number" id="module-${i}-y" step="0.1" value="0" required>
        </div>
        `;
        container.appendChild(moduleFieldset);
    }
}

function createModuleDisplays(robot) {
    const grid = document.getElementById('module-grid');
    grid.innerHTML = ''; // Delete any pre-existing elements before creating new ones

    const modules = robot.modules;
    modules.forEach((module, i) => {
        const article = document.createElement('article');
        article.className = 'module-display';
        const name = module.name;

        article.innerHTML = `
            <h3>${name}</h3>
            <div class="readout">
                <span class="label">Angle:</span>
                <span id="module-${i}-angle" class="value">0.0°</span>
            </div>
            <div class="readout">
                <span class="label">Speed:</span>
                <span id="module-${i}-speed" class="value">0.00 pixels/s</span>
            </div>
        `;
        grid.appendChild(article);
    });
}

function updateModuleDisplays(robot) {
    const configName = document.getElementById('config-name');
    configName.textContent = robot.name;
    const moduleCount = document.getElementById('module-count-display');
    moduleCount.textContent = robot.modules.length;

    // Update gyro heading display
    const gyroHeadingDisplay = document.getElementById('gyro-heading-display');
    if (gyroHeadingDisplay) {
        const headingDeg = (robot.gyroHeading * 180 / Math.PI).toFixed(1);
        gyroHeadingDisplay.textContent = `${headingDeg}°`;
    }

    const modules = robot.modules;
    modules.forEach((module, i) => {
        const angleElement = document.getElementById(`module-${i}-angle`);
        const speedElement = document.getElementById(`module-${i}-speed`);

        if (angleElement && speedElement) {
            const angleDeg = (module.angle * 180 / Math.PI).toFixed(1);
            angleElement.textContent = `${angleDeg}°`;
            speedElement.textContent = `${module.speed.toFixed(2)} pixels/s`;
        }
    });
}

/*
* END DYNAMIC DOM FUNCTIONS
* BEGIN ANIMATION CODE
*/

// Get the canvas and context as constants
const canvas = document.getElementById('swerve-canvas');
const ctx = canvas.getContext('2d');

// Get CSS variables for use in canvas
const rootStyles = getComputedStyle(document.documentElement);

function drawGrid(ctx, sideLength, gridSquareSize, xOffset, yOffset) {
    ctx.save();

    ctx.strokeStyle = rootStyles.getPropertyValue('--grid-color');
    ctx.lineWidth = 1;
    const startX = (-sideLength / 2) - xOffset;
    const endX = (sideLength / 2) - xOffset;
    const startY = (-sideLength / 2) - yOffset;
    const endY = (sideLength / 2) - yOffset;

    // Draw vertical lines
    for (let i = startX; i <= endX; i += gridSquareSize) {
        ctx.beginPath();
        ctx.moveTo(i, -sideLength / 2);
        ctx.lineTo(i, sideLength / 2);
        ctx.stroke();
    }

    // Draw horizontal lines
    for (let i = startY; i <= endY; i += gridSquareSize) {
        ctx.beginPath();
        ctx.moveTo(-sideLength / 2, i);
        ctx.lineTo(sideLength / 2, i);
        ctx.stroke();
    }

    ctx.restore();
}

function drawModule(ctx, module) {
    const x = module.position.x;
    const y = module.position.y;
    const arrowLength = Math.max(module.speed / 2, 5);

    ctx.save();
    ctx.translate(x, y);

    ctx.fillStyle = rootStyles.getPropertyValue('--swerve-fill-color');
    ctx.beginPath();
    ctx.arc(0, 0, 10, 0, Math.PI * 2);
    ctx.fill();

    ctx.strokeStyle = rootStyles.getPropertyValue('--swerve-module-color');
    ctx.lineWidth = 4;
    ctx.stroke();

    // Draw velocity arrow if module is moving
    if (module.speed > 0.01) {
        ctx.strokeStyle = rootStyles.getPropertyValue('--swerve-arrow-color');
        ctx.fillStyle = rootStyles.getPropertyValue('--swerve-arrow-color');
        ctx.lineWidth = 4;

        const endX = arrowLength * Math.cos(module.angle);
        const endY = arrowLength * Math.sin(module.angle);

        // Arrow line
        ctx.beginPath();
        ctx.moveTo(0, 0);
        ctx.lineTo(endX, endY);
        ctx.stroke();
    }


    ctx.restore();
}

function drawRobot(ctx, robot, heading) {
    ctx.save(); // Save current state before rotation

    ctx.rotate(heading);

    ctx.strokeStyle = rootStyles.getPropertyValue('--robot-frame-color')
    ctx.fillStyle = rootStyles.getPropertyValue('--robot-fill-color');
    ctx.lineWidth = 4;

    const modules = robot.modules.sort((a, b) => Math.atan2(a.position.y, a.position.x) - Math.atan2(b.position.y, b.position.x));

    ctx.beginPath();
    ctx.moveTo(modules[0].position.x, modules[0].position.y);
    for (let i = 1; i < modules.length; i++) {
        ctx.lineTo(modules[i].position.x, modules[i].position.y);
    }
    ctx.closePath();
    ctx.fill();
    ctx.stroke();

    modules.forEach(module => drawModule(ctx, module));

    ctx.restore(); // Restore to remove rotation
}


// Initialize Variables
const robotSize = 200;
const defaultModulePositions = PresetConfigs.fourWheelSquare(robotSize);
const robot = new SwerveDrive(defaultModulePositions, "4-Wheel Square");
createModuleDisplays(robot);
let xSpeed = 0;
let ySpeed = 0;
let turnSpeed = -1;

let gridSquareSize = 50;
let xGridOffset = 0;
let yGridOffset = 0;
robot.drive(xSpeed, ySpeed, 0, 500);

function animate() {
    // Clear and set up canvas
    ctx.clearRect(0, 0, canvas.width, canvas.height);
    ctx.save();
    ctx.translate(canvas.width / 2, canvas.height / 2);

    // Update speeds based on sliders
    xSpeed = parseFloat(vxSlider.value);
    ySpeed = -parseFloat(vySlider.value);
    turnSpeed = parseFloat(omegaSlider.value);

    // Animate the grid with robot movement
    let offsetSpeedDivisor = (100 - gridSquareSize <= 0 ? 1 : 100 - gridSquareSize);

    // Update grid offsets based on robot movement
    xGridOffset = (xGridOffset + (xSpeed / offsetSpeedDivisor)) % gridSquareSize;
    yGridOffset = (yGridOffset + (ySpeed / offsetSpeedDivisor)) % gridSquareSize;

    // Update module states before drawing the robot
    // The drive() method will update the gyroHeading internally
    robot.drive(xSpeed, ySpeed, turnSpeed, parseFloat(maxSpeedSlider.value));
    updateModuleDisplays(robot);

    // Draw the robot and it's movement. Grid should be oversized so movement
    // doesn't find the edge of the grid
    drawGrid(ctx, canvas.width * 2, gridSquareSize, xGridOffset, yGridOffset);
    drawRobot(ctx, robot, robot.gyroHeading);

    // Do it all over again
    ctx.restore();
    requestAnimationFrame(animate);
}

animate();