world_generator_tool/planet_visualization.js

// Global variables
let planetData = null;
let threejsScene, threejsCamera, threejsRenderer, threejsControls, threejsSphere;
let d3Svg, d3Projection, d3Path;
let progressInterval;

// Screen management
function showScreen(screenId) {
    // Hide all screens
    document.querySelectorAll('.screen').forEach(screen => {
        screen.classList.remove('active');
    });
    
    // Show the requested screen
    document.getElementById(screenId).classList.add('active');
    
    // Initialize visualization if showing that screen
    if (screenId === 'visualization-screen' && planetData) {
        // Need a small delay for the containers to be visible with correct dimensions
        setTimeout(() => {
            if (!threejsScene) {
                initVisualizations();
            } else {
                // If already initialized, handle resize for containers that were hidden
                handleResize();
            }
        }, 100);
    }
}

// Handle window resize
function handleResize() {
    if (threejsRenderer && threejsCamera) {
        const container = document.getElementById('threejs-container');
        const width = container.clientWidth;
        const height = container.clientHeight;
        
        threejsCamera.aspect = width / height;
        threejsCamera.updateProjectionMatrix();
        threejsRenderer.setSize(width, height);
    }
    
    if (d3Svg && d3Projection) {
        const container = document.getElementById('d3-container');
        const width = container.clientWidth;
        const height = container.clientHeight;
        
        d3Svg.attr('width', width)
            .attr('height', height);
        
        d3Projection.scale(width / (2 * Math.PI) * 0.9)
            .translate([width / 2, height / 1.8]);
        
        updateVisualization();
    }
}

// Initialize the 3D visualization with Three.js
function init3DVisualization() {
    const container = document.getElementById('threejs-container');
    const width = container.clientWidth;
    const height = container.clientHeight;
    
    // Create scene
    threejsScene = new THREE.Scene();
    threejsScene.background = new THREE.Color(0x000000);
    
    // Create camera
    threejsCamera = new THREE.PerspectiveCamera(75, width / height, 0.1, 1000);
    threejsCamera.position.z = 2;
    
    // Add polar axis helper
    const axisHelper = new THREE.Group();
    
    // North-South pole axis (red)
    const poleGeometry = new THREE.CylinderGeometry(0.01, 0.01, 2.4, 8);
    const poleMaterial = new THREE.MeshBasicMaterial({ color: 0xff0000 });
    const poleAxis = new THREE.Mesh(poleGeometry, poleMaterial);
    // Correctly align with y-axis (poles should be on y-axis)
    poleAxis.rotation.z = 0;
    axisHelper.add(poleAxis);
    
    // North pole cap (red)
    const northCapGeometry = new THREE.ConeGeometry(0.04, 0.1, 8);
    const northCap = new THREE.Mesh(northCapGeometry, poleMaterial);
    northCap.position.set(0, 1.25, 0);
    northCap.rotation.x = Math.PI; // Point up
    axisHelper.add(northCap);
    
    // South pole cap (red)
    const southCapGeometry = new THREE.ConeGeometry(0.04, 0.1, 8);
    const southCap = new THREE.Mesh(southCapGeometry, poleMaterial);
    southCap.position.set(0, -1.25, 0);
    axisHelper.add(southCap);
    
    threejsScene.add(axisHelper);
    
    // Create renderer
    threejsRenderer = new THREE.WebGLRenderer({ antialias: true });
    threejsRenderer.setSize(width, height);
    container.appendChild(threejsRenderer.domElement);
    
    // Add orbit controls
    threejsControls = new THREE.OrbitControls(threejsCamera, threejsRenderer.domElement);
    threejsControls.enableDamping = true;
    threejsControls.dampingFactor = 0.05;
    
    // Add ambient light
    const ambientLight = new THREE.AmbientLight(0xffffff, 0.5);
    threejsScene.add(ambientLight);
    
    // Add directional light
    const directionalLight = new THREE.DirectionalLight(0xffffff, 0.8);
    directionalLight.position.set(1, 1, 1);
    threejsScene.add(directionalLight);
    
    // Start animation loop
    animate();
}

// Initialize the 2D visualization with D3.js
function init2DVisualization() {
    const container = document.getElementById('d3-container');
    const width = container.clientWidth;
    const height = container.clientHeight;
    
    // Create SVG
    d3Svg = d3.select('#d3-container')
        .append('svg')
        .attr('width', width)
        .attr('height', height);
    
    // Create projection
    d3Projection = d3.geoMercator()
        .scale(width / (2 * Math.PI) * 0.9)
        .translate([width / 2, height / 2]);
    
    // Create path generator
    d3Path = d3.geoPath()
        .projection(d3Projection);
    
    // Add a background rectangle representing the ocean
    d3Svg.append('rect')
        .attr('width', width)
        .attr('height', height)
        .attr('fill', '#a4d1e9');
    
    // Load and draw actual world map data
    d3.json('https://cdn.jsdelivr.net/npm/world-atlas@2/countries-110m.json').then(world => {
        // Draw countries
        const countries = topojson.feature(world, world.objects.countries);
        
        d3Svg.append('g')
            .attr('class', 'countries')
            .selectAll('path')
            .data(countries.features)
            .enter()
            .append('path')
            .attr('d', d3Path)
            .attr('fill', '#d3b683')
            .attr('stroke', '#a89070')
            .attr('stroke-width', 0.5)
            .attr('opacity', 0.8);
            
        // Add graticule (grid lines)
        const graticule = d3.geoGraticule()
            .step([15, 15]);
        
        d3Svg.append('path')
            .datum(graticule)
            .attr('class', 'graticule')
            .attr('d', d3Path)
            .style('fill', 'none')
            .style('stroke', '#ffffff')
            .style('stroke-width', '0.5px')
            .style('opacity', 0.3);
        
        // Draw equator
        const equator = {
            type: "LineString",
            coordinates: [[-180, 0], [-90, 0], [0, 0], [90, 0], [180, 0]]
        };
        
        d3Svg.append('path')
            .datum(equator)
            .attr('class', 'equator')
            .attr('d', d3Path)
            .style('fill', 'none')
            .style('stroke', '#00ffff')
            .style('stroke-width', '2px');
    });
}

// Create a Three.js sphere from points
function createThreeJSSphere(points) {
    // Remove existing sphere if any
    if (threejsSphere) {
        threejsScene.remove(threejsSphere);
    }
    
    // Create a group to hold the sphere and equator
    threejsSphere = new THREE.Group();
    
    // Create geometry
    const geometry = new THREE.SphereGeometry(1, 32, 32);
    
    // Create material
    const material = new THREE.MeshPhongMaterial({
        vertexColors: true,
        flatShading: true
    });
    
    // Add equator
    const equatorGeometry = new THREE.TorusGeometry(1.01, 0.005, 16, 100);
    const equatorMaterial = new THREE.MeshBasicMaterial({ color: 0x00ffff });
    const equator = new THREE.Mesh(equatorGeometry, equatorMaterial);
    equator.rotation.x = Math.PI / 2; // Rotate to lie on the x-z plane
    threejsSphere.add(equator);
    
    // Add colors to geometry
    const colors = [];
    const positionAttribute = geometry.getAttribute('position');
    
    // For each vertex in the sphere
    for (let i = 0; i < positionAttribute.count; i++) {
        const vertex = new THREE.Vector3();
        vertex.fromBufferAttribute(positionAttribute, i);
        vertex.normalize();
        
        // Find closest point in our dataset
        let minDist = Infinity;
        let closestPoint = null;
        
        for (const point of points) {
            const pointVec = new THREE.Vector3(point.x, point.y, point.z);
            const dist = vertex.distanceTo(pointVec);
            if (dist < minDist) {
                minDist = dist;
                closestPoint = point;
            }
        }
        
        // Set color from closest point
        colors.push(
            closestPoint.color[0],
            closestPoint.color[1],
            closestPoint.color[2]
        );
    }
    
    // Add colors to geometry
    geometry.setAttribute('color', new THREE.Float32BufferAttribute(colors, 3));
    
    // Create mesh for the sphere
    const sphereMesh = new THREE.Mesh(geometry, material);
    
    // Add sphere to the group
    threejsSphere.add(sphereMesh);
    
    // Add group to scene
    threejsScene.add(threejsSphere);
}

// Create a D3.js map from points
function createD3Map(points) {
    // Clear existing points but preserve graticule, equator, and axis
    d3Svg.selectAll('.planet-point').remove();
    
    // Create groups for continents to improve organization
    const pointsGroup = d3Svg.append('g').attr('class', 'points-group');
    
    // Plot each point on the map
    points.forEach(point => {
        // Convert spherical coordinates to longitude/latitude
        // phi goes from -π to π, convert to -180 to 180 degrees (longitude)
        const lon = (point.phi * 180 / Math.PI);
        
        // theta goes from 0 to π, convert to 90 to -90 degrees (latitude)
        const lat = 90 - (point.theta * 180 / Math.PI);
        
        // Skip points that might cause projection issues at the poles
        if (lat > 85 || lat < -85) return;
        
        // Only draw the point if it projects properly
        const projected = d3Projection([lon, lat]);
        if (projected) {
            // Calculate scale factor based on latitude to mimic Mercator distortion
            // 1/cos(latitude in radians) gives us the characteristic Mercator scaling
            const latRad = lat * Math.PI / 180;
            const scaleFactor = Math.min(3, 1 / Math.cos(latRad)); // Cap at 3x to avoid extreme sizes
            
            pointsGroup.append('circle')
                .attr('class', 'planet-point')
                .attr('cx', projected[0])
                .attr('cy', projected[1])
                .attr('r', 1.5 * scaleFactor) // Scale radius by the Mercator factor
                .style('fill', `rgb(${point.color[0] * 255}, ${point.color[1] * 255}, ${point.color[2] * 255})`)
                .style('opacity', 0.8);
        }
    });
}

// Animation loop for Three.js
function animate() {
    requestAnimationFrame(animate);
    
    // Update controls
    if (threejsControls) {
        threejsControls.update();
    }
    
    // Auto-rotate if enabled and on visualization screen
    if (document.getElementById('visualization-screen').classList.contains('active') && 
        document.getElementById('auto-rotate').checked && 
        threejsSphere) {
        threejsSphere.rotation.y += 0.005;
    }
    
    // Render scene
    if (threejsRenderer && threejsScene && threejsCamera) {
        threejsRenderer.render(threejsScene, threejsCamera);
    }
}

// Update visualizations with new data
function updateVisualization() {
    if (planetData && planetData.points) {
        createThreeJSSphere(planetData.points);
        createD3Map(planetData.points);
    }
}

// Initialize visualizations
function initVisualizations() {
    // Only initialize if not already done
    if (!threejsScene) {
        init3DVisualization();
    }
    
    if (!d3Svg) {
        init2DVisualization();
    }
    
    updateVisualization();
}

// Start progress bar simulation
function startProgressBar() {
    let progress = 0;
    const progressBar = document.getElementById('progress-bar');
    const progressText = document.getElementById('progress-text');
    
    clearInterval(progressInterval);
    
    progressBar.style.width = '0%';
    progressText.textContent = '0%';
    
    // Simulate progress updates
    progressInterval = setInterval(() => {
        if (progress >= 100) {
            clearInterval(progressInterval);
            return;
        }
        
        // Make progress increases less predictable
        const increment = Math.random() * 5 + 1;
        progress = Math.min(progress + increment, 100);
        
        progressBar.style.width = `${progress}%`;
        progressText.textContent = `${Math.round(progress)}%`;
        
        // When complete, show visualization
        if (progress >= 100) {
            setTimeout(() => {
                showScreen('visualization-screen');
            }, 500); // Short delay for animation
        }
    }, 100);
}

// Generate planet
function generatePlanet() {
    const pointsCount = parseInt(document.getElementById('points-count').value);
    const colorScheme = document.getElementById('color-scheme').value;
    
    // Show loading screen
    showScreen('loading-screen');
    
    // Start progress bar
    startProgressBar();
    
    // Call Python function to generate data with a slight delay
    // to allow the progress bar to start
    setTimeout(() => {
        pywebview.api.regenerate_planet(pointsCount, colorScheme)
            .then(data => {
                planetData = data;
                
                // Ensure we're at 100% when done
                document.getElementById('progress-bar').style.width = '100%';
                document.getElementById('progress-text').textContent = '100%';
                
                // Show visualization after a short delay
                setTimeout(() => {
                    showScreen('visualization-screen');
                }, 300);
            })
            .catch(error => {
                console.error('Error generating planet:', error);
                alert('An error occurred while generating the planet. Please try again.');
                showScreen('settings-screen');
            });
    }, 500);
}

// Set planet data and update visualizations
function setPlanetData(data) {
    planetData = data;
    
    // Ensure progress is complete and move to visualization
    document.getElementById('progress-bar').style.width = '100%';
    document.getElementById('progress-text').textContent = '100%';
    
    // Clear any progress simulation
    clearInterval(progressInterval);
    
    // Show visualization screen after a short delay
    setTimeout(() => {
        showScreen('visualization-screen');
    }, 300);
}

// Set up event listeners
function setupEventListeners() {
    // Generate button on settings screen
    document.getElementById('generate-btn').addEventListener('click', generatePlanet);
    
    // Regenerate button on visualization screen
    document.getElementById('regenerate-btn').addEventListener('click', () => {
        const pointsCount = parseInt(document.getElementById('points-count').value);
        const colorScheme = document.getElementById('color-scheme').value;
        
        // Show loading screen
        showScreen('loading-screen');
        
        // Start progress bar
        startProgressBar();
        
        // Call Python function to regenerate data
        setTimeout(() => {
            pywebview.api.regenerate_planet(pointsCount, colorScheme)
                .then(data => {
                    planetData = data;
                    setTimeout(() => {
                        showScreen('visualization-screen');
                    }, 300);
                });
        }, 500);
    });
    
    // Back button on visualization screen
    document.getElementById('back-btn').addEventListener('click', () => {
        showScreen('settings-screen');
    });
    
    // Window resize handler
    window.addEventListener('resize', handleResize);
}

// Initialize when document is ready
document.addEventListener('DOMContentLoaded', () => {
    setupEventListeners();
    
    // Start on the settings screen
    showScreen('settings-screen');
    
    // Start the animation loop anyway for when we need it
    animate();
});