/**
 * File: cloth.component.ts
 * Description: Component for cloth simulation using WebGPU compute shaders.
 */

import { Component } from '@angular/core';
import { MatCard, MatCardContent, MatCardHeader, MatCardTitle } from '@angular/material/card';
import { TranslatePipe } from '@ngx-translate/core';
import { BabylonCanvas, RenderConfig, SceneEventData } from '../../../shared/components/render-canvas/babylon-canvas.component';
import {ComputeShader, StorageBuffer, MeshBuilder, ShaderMaterial, ShaderLanguage, ArcRotateCamera, GroundMesh, WebGPUEngine, Scene} from '@babylonjs/core';
import {
  CLOTH_FRAGMENT_SHADER_WGSL,
  CLOTH_INTEGRATE_COMPUTE_WGSL,
  CLOTH_SOLVE_COMPUTE_WGSL,
  CLOTH_VELOCITY_COMPUTE_WGSL,
  CLOTH_VERTEX_SHADER_WGSL
} from './cloth.shader';
import {MatButton} from '@angular/material/button';
import {ClothBuffers, ClothConfig, ClothData, ClothPipelines} from './cloth.model';
import {Information} from '../information/information';
import {AlgorithmInformation} from '../information/information.models';
import {UrlConstants} from '../../../constants/UrlConstants';

@Component({
  selector: 'app-cloth',
  imports: [
    MatCard,
    MatCardContent,
    MatCardHeader,
    MatCardTitle,
    TranslatePipe,
    BabylonCanvas,
    MatButton,
    Information
  ],
  templateUrl: './cloth.component.html',
  styleUrl: './cloth.component.scss',
})
export class ClothComponent {
  private currentSceneData: SceneEventData | null = null;
  private simulationTime: number = 0;
  private clothMesh: GroundMesh | null = null;
  public isWindActive: boolean = false;
  public isOutlineActive: boolean = false;

  public renderConfig: RenderConfig = {
    mode: '3D',
    initialViewSize: 20,
    shaderLanguage: ShaderLanguage.WGSL
  };

  algoInformation: AlgorithmInformation = {
    title: 'CLOTH.EXPLANATION.TITLE',
    entries: [
      {
        name: 'CLOTH.EXPLANATION.CLOTH_SIMULATION_EXPLANATION_TITLE',
        description: 'CLOTH.EXPLANATION.CLOTH_SIMULATION_EXPLANATION',
        link: UrlConstants.CLOTH_SIMULATION_WIKI,
        translateName: true
      },
      {
        name: 'CLOTH.EXPLANATION.XPBD_EXPLANATION_TITLE',
        description: 'CLOTH.EXPLANATION.XPBD_EXPLANATION',
        link: UrlConstants.XPBD_WIKI,
        translateName: true
      },
      {
        name: 'CLOTH.EXPLANATION.GPU_PARALLELIZATION_EXPLANATION_TITLE',
        description: 'CLOTH.EXPLANATION.GPU_PARALLELIZATION_EXPLANATION',
        link: UrlConstants.GPU_COMPUTING_WIKI,
        translateName: true
      },
      {
        name: 'CLOTH.EXPLANATION.DATA_STRUCTURES_EXPLANATION_TITLE',
        description: 'CLOTH.EXPLANATION.DATA_STRUCTURES_EXPLANATION',
        link: UrlConstants.DATA_STRUCTURE_WIKI,
        translateName: true
      }
    ],
    disclaimer: 'CLOTH.EXPLANATION.DISCLAIMER',
    disclaimerBottom: '',
    disclaimerListEntry: ['CLOTH.EXPLANATION.DISCLAIMER_1', 'CLOTH.EXPLANATION.DISCLAIMER_2', 'CLOTH.EXPLANATION.DISCLAIMER_3', 'CLOTH.EXPLANATION.DISCLAIMER_4']
  };

  /**
   * Called when the Babylon scene is ready.
   * @param event The scene event data.
   */
  public onSceneReady(event: SceneEventData): void {
    this.currentSceneData = event;
    this.createSimulation();
  }

  public toggleWind(): void {
    this.isWindActive = !this.isWindActive;
  }

  public toggleMesh(): void {
    this.isOutlineActive = !this.isOutlineActive;
    if (!this.clothMesh?.material) {
      return;
    }
    this.clothMesh.material.wireframe = this.isOutlineActive;
  }

  /**
   * Initializes and starts the cloth simulation.
   */
  private createSimulation(): void {
    if (!this.currentSceneData) return;

    const { engine, scene } = this.currentSceneData;

    // 1. Define physics parameters
    const config = this.getClothConfig();

    // 2. Generate initial CPU data (positions, constraints)
    const clothData = this.generateClothData(config);

    // 3. Upload to GPU
    const buffers = this.createStorageBuffers(engine, clothData);

    // 4. Create Compute Shaders
    const pipelines = this.setupComputePipelines(engine, buffers);

    // 5. Setup Rendering (Mesh, Material, Camera)
    this.setupRenderMesh(scene, config, buffers.positions);

    // 6. Start the physics loop
    this.startRenderLoop(engine, scene, config, buffers, pipelines);
  }

  // ========================================================================
  // 1. CONFIGURATION
  // ========================================================================
  private getClothConfig(): ClothConfig {
    const gridWidth = 100;
    const gridHeight = 100;
    const spacing = 0.05;
    const density = 1.0;
    const particleArea = spacing * spacing;
    const particleMass = density * particleArea;

    return {
      gridWidth,
      gridHeight,
      spacing,
      density,
      numVertices: gridWidth * gridHeight,
      particleInvMass: 1.0 / particleMass
    };
  }

  // ========================================================================
  // 2. DATA GENERATION (CPU)
  // ========================================================================
  private generateClothData(config: ClothConfig): ClothData {
    const positionsData = new Float32Array(config.numVertices * 4);
    const prevPositionsData = new Float32Array(config.numVertices * 4);
    const velocitiesData = new Float32Array(config.numVertices * 4);

    const constraintsP0: number[] = [];
    const constraintsP1: number[] = [];
    const constraintsP2: number[] = [];
    const constraintsP3: number[] = [];

    const addConstraint = (arr: number[], a: number, b: number): void => {
      arr.push(a, b, config.spacing, 1.0);
    };

    // Fill positions (Pin top row)
    for (let y = 0; y < config.gridHeight; y++) {
      for (let x = 0; x < config.gridWidth; x++) {
        const idx = (y * config.gridWidth + x) * 4;
        positionsData[idx + 0] = (x - config.gridWidth / 2) * config.spacing;
        positionsData[idx + 1] = 5.0 - (y * config.spacing);
        positionsData[idx + 2] = 0.0;
        positionsData[idx + 3] = (y === 0) ? 0.0 : config.particleInvMass;

        prevPositionsData[idx + 0] = positionsData[idx + 0];
        prevPositionsData[idx + 1] = positionsData[idx + 1];
        prevPositionsData[idx + 2] = positionsData[idx + 2];
        prevPositionsData[idx + 3] = positionsData[idx + 3];
      }
    }

    // Graph Coloring (4 Phases)
    for (let y = 0; y < config.gridHeight; y++) {
      for (let x = 0; x < config.gridWidth - 1; x += 2) addConstraint(constraintsP0, y * config.gridWidth + x, y * config.gridWidth + x + 1);
      for (let x = 1; x < config.gridWidth - 1; x += 2) addConstraint(constraintsP1, y * config.gridWidth + x, y * config.gridWidth + x + 1);
    }
    for (let y = 0; y < config.gridHeight - 1; y += 2) {
      for (let x = 0; x < config.gridWidth; x++) addConstraint(constraintsP2, y * config.gridWidth + x, (y + 1) * config.gridWidth + x);
    }
    for (let y = 1; y < config.gridHeight - 1; y += 2) {
      for (let x = 0; x < config.gridWidth; x++) addConstraint(constraintsP3, y * config.gridWidth + x, (y + 1) * config.gridWidth + x);
    }

    const constraintsP4: number[] = [];
    const constraintsP5: number[] = [];
    const constraintsP6: number[] = [];
    const constraintsP7: number[] = [];

    const diagSpacing = config.spacing * Math.SQRT2;
    const addDiagConstraint = (arr: number[], a: number, b: number): void => {
      arr.push(a, b, diagSpacing, 1.0);
    };

    for (let y = 0; y < config.gridHeight - 1; y++) {
      const arr = (y % 2 === 0) ? constraintsP4 : constraintsP5;
      for (let x = 0; x < config.gridWidth - 1; x++) {
        addDiagConstraint(arr, y * config.gridWidth + x, (y + 1) * config.gridWidth + (x + 1));
      }
    }

    for (let y = 0; y < config.gridHeight - 1; y++) {
      const arr = (y % 2 === 0) ? constraintsP6 : constraintsP7;
      for (let x = 0; x < config.gridWidth - 1; x++) {
        addDiagConstraint(arr, y * config.gridWidth + (x + 1), (y + 1) * config.gridWidth + x);
      }
    }

    return {
      positions: positionsData,
      prevPositions: prevPositionsData,
      velocities: velocitiesData,
      constraints: [
        constraintsP0, constraintsP1, constraintsP2, constraintsP3,
        constraintsP4, constraintsP5, constraintsP6, constraintsP7
      ],
      params: new Float32Array(8)
    };
  }

  // ========================================================================
  // 3. BUFFER CREATION (GPU)
  // ========================================================================
  private createStorageBuffers(engine: WebGPUEngine, data: ClothData): ClothBuffers {
    const createBuffer = (arrayData: Float32Array | number[]): StorageBuffer => {
      const buffer = new StorageBuffer(engine, arrayData.length * 4);
      buffer.update(arrayData instanceof Float32Array ? arrayData : new Float32Array(arrayData));
      return buffer;
    };

    return {
      positions: createBuffer(data.positions),
      prevPositions: createBuffer(data.prevPositions),
      velocities: createBuffer(data.velocities),
      params: createBuffer(data.params),
      constraints: data.constraints.map(cData => createBuffer(cData))
    };
  }

  // ========================================================================
  // 4. COMPUTE SHADERS
  // ========================================================================
  private setupComputePipelines(engine: WebGPUEngine, buffers: ClothBuffers): ClothPipelines {

    // Helper for integrating & velocity
    const createBasicShader = (name: string, source: string) => {
      const cs = new ComputeShader(name, engine, { computeSource: source }, {
        bindingsMapping: {
          "p": { group: 0, binding: 0 },
          "positions": { group: 0, binding: 1 },
          "prev_positions": { group: 0, binding: 2 },
          "velocities": { group: 0, binding: 3 }
        }
      });
      cs.setStorageBuffer("p", buffers.params);
      cs.setStorageBuffer("positions", buffers.positions);
      cs.setStorageBuffer("prev_positions", buffers.prevPositions);
      cs.setStorageBuffer("velocities", buffers.velocities);
      return cs;
    };

    // Helper for solvers
    const createSolverShader = (name: string, constraintBuffer: StorageBuffer) => {
      const cs = new ComputeShader(name, engine, { computeSource: CLOTH_SOLVE_COMPUTE_WGSL }, {
        bindingsMapping: {
          "p": { group: 0, binding: 0 },
          "positions": { group: 0, binding: 1 },
          "constraints": { group: 0, binding: 2 }
        }
      });
      cs.setStorageBuffer("p", buffers.params);
      cs.setStorageBuffer("positions", buffers.positions);
      cs.setStorageBuffer("constraints", constraintBuffer);
      return cs;
    };

    return {
      integrate: createBasicShader("integrate", CLOTH_INTEGRATE_COMPUTE_WGSL),
      solvers: buffers.constraints.map((cBuffer, i) => createSolverShader(`solve${i}`, cBuffer)),
      velocity: createBasicShader("velocity", CLOTH_VELOCITY_COMPUTE_WGSL)
    };
  }

  // ========================================================================
  // 5. RENDERING SETUP
  // ========================================================================
  private setupRenderMesh(scene: Scene, config: ClothConfig, positionsBuffer: StorageBuffer): void {
    if (this.clothMesh) {
      scene.removeMesh(this.clothMesh);
    }

    this.clothMesh = MeshBuilder.CreateGround("cloth", { width: 10, height: 10, subdivisions: config.gridWidth - 1 }, scene);

    const clothMaterial = new ShaderMaterial("clothMat", scene, {
      vertexSource: CLOTH_VERTEX_SHADER_WGSL,
      fragmentSource: CLOTH_FRAGMENT_SHADER_WGSL
    }, {
      attributes: ["position", "uv"],
      uniforms: ["viewProjection"],
      storageBuffers: ["positions"],
      shaderLanguage: ShaderLanguage.WGSL
    });

    clothMaterial.backFaceCulling = false;
    clothMaterial.setStorageBuffer("positions", positionsBuffer);
    this.clothMesh.material = clothMaterial;

    const camera = scene.activeCamera as ArcRotateCamera;
    if (camera) {
      camera.alpha = Math.PI / 4;
      camera.beta = Math.PI / 2.5;
      camera.radius = 15;
    }
  }

  // ========================================================================
  // 6. RENDER LOOP
  // ========================================================================
  private startRenderLoop(engine: WebGPUEngine, scene: Scene, config: ClothConfig, buffers: ClothBuffers, pipelines: ClothPipelines): void {
    const paramsData = new Float32Array(8);

    // Pre-calculate constraint dispatch sizes for the 4 phases
    const constraintsLength = buffers.constraints.map(b => (b as any)._buffer.capacity / 4 / 4); // Elements / vec4 length
    const dispatchXConstraints = constraintsLength.map(len => Math.ceil(len / 64));
    const dispatchXVertices = Math.ceil(config.numVertices / 64);
    const substeps = 15;

    scene.onBeforeRenderObservable.clear();
    scene.onBeforeRenderObservable.add(() => {
      this.simulationTime += engine.getDeltaTime() / 1000.0;

      // Update Physics Parameters
      const windX = this.isWindActive ? 5.0 : 0.0;
      const windY = 0.0;
      const windZ = this.isWindActive ? 15.0 : 0.0;
      const scaledCompliance = 0.00001 * config.particleInvMass * config.spacing;

      paramsData[0] = 0.016; // dt
      paramsData[1] = -9.81; // gravity
      paramsData[2] = scaledCompliance;
      paramsData[3] = config.numVertices;
      paramsData[4] = windX;
      paramsData[5] = windY;
      paramsData[6] = windZ;
      paramsData[7] = this.simulationTime;

      buffers.params.update(paramsData);

      // 1. Predict positions
      pipelines.integrate.dispatch(dispatchXVertices, 1, 1);

      // 2. XPBD Solver (Substeps) - Graph Coloring Phase
      for (let i = 0; i < substeps; i++) {
        for (let phase = 0; phase < pipelines.solvers.length; phase++) {
          pipelines.solvers[phase].dispatch(dispatchXConstraints[phase], 1, 1);
        }
      }

      // 3. Update velocities
      pipelines.velocity.dispatch(dispatchXVertices, 1, 1);
    });
  }
}
