// 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) {
        // Take the requested speed and turn rate of the robot and calculate
        // speed and angle of this module to achieve it

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

        // Combine translation and rotation
        this.velocity.x = velocityX + rotX;
        this.velocity.y = velocityY + rotY;

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

// Swerve drive class to represent the robot as a whole
class SwerveDrive {
    constructor(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)
        );
    }

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

        // 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 > maxSpeed) {
            const scale = maxSpeed / maxCalculated;
            this.modules.forEach(module => {
                module.speed *= scale;
            });
        }
    }
}

// Preset robot generators
const PresetConfigs = {
    twoWheel: (size) => [
        { x: 0, y: size / 2, name: "Left" },
        { x: 0, y: -size / 2, 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.7;
        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;
    }
};

// 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 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');

// Add event listeners for drive controls
vxSlider.addEventListener('input', (e) => {
    vxOutput.textContent = parseFloat(e.target.value).toFixed(1);
});

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

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

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

// 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) {
    ctx.strokeStyle = rootStyles.getPropertyValue('--grid-color');
    ctx.lineWidth = 1;

    for (let i = -sideLength / 2; i <= sideLength / 2; i += gridSquareSize) {
        ctx.beginPath();
        ctx.moveTo(i, -sideLength / 2);
        ctx.lineTo(i, sideLength / 2);
        ctx.stroke();

        ctx.beginPath();
        ctx.moveTo(-sideLength / 2, i);
        ctx.lineTo(sideLength / 2, i);
        ctx.stroke();
    }
}

function drawRobot(ctx, robot) {
    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));
    console.log(modules);

    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();
}

ctx.translate(canvas.width / 2, canvas.height / 2);
drawGrid(ctx, 600, 50);
drawRobot(ctx, new SwerveDrive(PresetConfigs.eightWheel(200)));