using UnityEngine;
using System.Collections.Generic;
using System.Linq;
using System.IO; // Added for file writing
using System; // For Math.Sqrt

// NBodySimulator.cs
// This script manages the overall N-body gravitational simulation using Velocity Verlet integration.
// It instantiates all celestial bodies, calculates gravitational forces between them
// using real-world constants, updates their states, and centralizes all simulation data logging.

public class NBodySimulator : MonoBehaviour
{
    [Tooltip("The actual gravitational constant G = 6.6743e-11 N(m/kg)^2. Do not change this for empirical accuracy.")]
    public const double gravitationalConstant = 6.6743e-11; // N(m/kg)^2 - actual G (double precision)

    [Tooltip("The time step (dt) for each physics update in seconds. Smaller values increase accuracy but decrease performance.")]
    public float simulationTimeStep_s = 60f; // Example: 60 seconds (1 minute) per physics step

    [Tooltip("Scaling factor to convert real-world meters into Unity units. " +
             "E.g., 1e9 means 1 Unity Unit = 1,000,000,000 meters (1 Gigameter). " +
             "Adjust this to make planets visible and to prevent them from being too far apart in Unity's view.")]
    public float distanceScale_m_per_unity_unit = 1e9f; // 1 Unity unit represents 1 billion meters

    [Header("Visual References")]
    [Tooltip("Assign the GameObject for the Sun.")]
    public GameObject sunGameObject;
    [Tooltip("Assign the GameObject for the Earth.")]
    public GameObject earthGameObject;
    [Tooltip("Assign the GameObject for the Moon.")]
    public GameObject moonGameObject;

    [Header("Simulation Logging")]
    [Tooltip("Enable to log simulation data to a CSV file.")]
    public bool enableDataLogging = false;
    [Tooltip("Name of the log file (e.g., nbody_simulation_log.csv).")]
    public string logFileName = "nbody_simulation_log.csv";

    private List<PlanetController> _celestialBodies;
    private StreamWriter _dataLogger;

    void Awake()
    {
        _celestialBodies = new List<PlanetController>();

        // Get PlanetController components from assigned GameObjects and initialize them
        if (sunGameObject != null)
        {
            PlanetController sun = sunGameObject.GetComponent<PlanetController>();
            if (sun != null) { _celestialBodies.Add(sun); sun.Initialize(); }
            else { Debug.LogError("Sun GameObject is missing PlanetController component."); }
        }
        if (earthGameObject != null)
        {
            PlanetController earth = earthGameObject.GetComponent<PlanetController>();
            if (earth != null) { _celestialBodies.Add(earth); earth.Initialize(); }
            else { Debug.LogError("Earth GameObject is missing PlanetController component."); }
        }
        if (moonGameObject != null)
        {
            PlanetController moon = moonGameObject.GetComponent<PlanetController>();
            if (moon != null) { _celestialBodies.Add(moon); moon.Initialize(); }
            else { Debug.LogError("Moon GameObject is missing PlanetController component."); }
        }

        // Initialize data logger if enabled
        if (enableDataLogging)
        {
            InitializeDataLogger();
        }
    }

    void Start()
    {
        // Set initial positions in Unity based on real-world initial positions and scaling
        foreach (var body in _celestialBodies)
        {
            body.transform.position = (Vector3)(body.currentPosition_m / distanceScale_m_per_unity_unit);
        }
    }

    void FixedUpdate()
    {
        // Step 1: Reset accelerations for all bodies
        foreach (var body in _celestialBodies)
        {
            body.ResetAcceleration();
        }

        // Step 2: Calculate gravitational forces and accumulate accelerations (a(t))
        CalculateAllGravitationalForces();

        // Step 3: Update positions (x(t+dt)) and half-step velocities (v(t) + 0.5*a(t)*dt)
        foreach (var body in _celestialBodies)
        {
            body.UpdateKinematics(simulationTimeStep_s);
        }

        // Step 4: Recalculate forces at new positions to get new accelerations (a(t+dt))
        // This requires temporarily updating Unity positions for accurate distance calculations
        foreach (var body in _celestialBodies)
        {
            body.transform.position = (Vector3)(body.currentPosition_m / distanceScale_m_per_unity_unit);
        }
        Vector3d earthNewAcceleration = Vector3d.zero; // Store Earth's new acceleration separately
        Vector3d moonNewAcceleration = Vector3d.zero; // Store Moon's new acceleration separately

        // Need to re-calculate forces based on the new positions
        foreach (var p1 in _celestialBodies)
        {
            foreach (var p2 in _celestialBodies)
            {
                if (p1 == p2) continue;

                Vector3d r_vec = p2.currentPosition_m - p1.currentPosition_m;
                double r_mag_sq = r_vec.sqrMagnitude;

                if (r_mag_sq < 1e-9) continue; // Avoid division by zero for coincident bodies

                double r_mag = Math.Sqrt(r_mag_sq);
                Vector3d r_norm = r_vec.normalized;

                // F = G * m1 * m2 / r^2
                double forceMagnitude = (gravitationalConstant * p1.mass * p2.mass) / r_mag_sq;
                Vector3d force_on_p1 = r_norm * forceMagnitude; // Force on p1 due to p2

                // a = F / m
                Vector3d acceleration_on_p1 = force_on_p1 / p1.mass;

                // Accumulate new acceleration for p1
                p1.currentAcceleration_m_per_s_sq += acceleration_on_p1;

                // Special handling for Earth and Moon to pass forces to TrenchSimulator
                if (p1.name == "Earth" && p2.name == "Sun")
                {
                    FindObjectOfType<TrenchInstabilitySimulator>()?.gravitationalForceFromSun_N += force_on_p1;
                }
                else if (p1.name == "Earth" && p2.name == "Moon")
                {
                    FindObjectOfType<TrenchInstabilitySimulator>()?.gravitationalForceFromMoon_N += force_on_p1;
                }
            }
        }

        // Step 5: Finalize velocities (v(t+dt)) using a(t) and a(t+dt)
        foreach (var body in _celestialBodies)
        {
            // The currentAcceleration_m_per_s_sq now holds a(t+dt) from the re-calculation above
            body.FinalizeVelocity(body.currentAcceleration_m_per_s_sq, simulationTimeStep_s);
        }

        // Step 6: Update Unity positions for rendering
        foreach (var body in _celestialBodies)
        {
            body.transform.position = (Vector3)(body.currentPosition_m / distanceScale_m_per_unity_unit);
        }

        // Step 7: Log data
        if (enableDataLogging)
        {
            LogSimulationData();
        }
    }


    void CalculateAllGravitationalForces()
    {
        // Reset forces on Earth from Sun/Moon for TrenchSimulator
        FindObjectOfType<TrenchInstabilitySimulator>()?.gravitationalForceFromSun_N = Vector3d.zero;
        FindObjectOfType<TrenchInstabilitySimulator>()?.gravitationalForceFromMoon_N = Vector3d.zero;

        // Calculate forces between all pairs of bodies
        for (int i = 0; i < _celestialBodies.Count; i++)
        {
            PlanetController p1 = _celestialBodies[i];
            for (int j = i + 1; j < _celestialBodies.Count; j++) // Avoid duplicate calculations and self-interaction
            {
                PlanetController p2 = _celestialBodies[j];

                Vector3d r_vec = p2.currentPosition_m - p1.currentPosition_m;
                double r_mag_sq = r_vec.sqrMagnitude;

                if (r_mag_sq < 1e-9) continue; // Avoid division by zero for coincident bodies

                double r_mag = Math.Sqrt(r_mag_sq);
                Vector3d r_norm = r_vec.normalized;

                // F = G * m1 * m2 / r^2
                double forceMagnitude = (gravitationalConstant * p1.mass * p2.mass) / r_mag_sq;
                Vector3d force_on_p1 = r_norm * forceMagnitude; // Force on p1 due to p2
                Vector3d force_on_p2 = -force_on_p1; // Newton's third law

                // a = F / m
                p1.currentAcceleration_m_per_s_sq += force_on_p1 / p1.mass;
                p2.currentAcceleration_m_per_s_sq += force_on_p2 / p2.mass;

                // Pass forces to TrenchInstabilitySimulator if Earth is involved
                if (p1.name == "Earth")
                {
                    if (p2.name == "Sun") FindObjectOfType<TrenchInstabilitySimulator>()?.gravitationalForceFromSun_N += force_on_p1;
                    if (p2.name == "Moon") FindObjectOfType<TrenchInstabilitySimulator>()?.gravitationalForceFromMoon_N += force_on_p1;
                }
                else if (p2.name == "Earth")
                {
                    if (p1.name == "Sun") FindObjectOfType<TrenchInstabilitySimulator>()?.gravitationalForceFromSun_N += force_on_p2;
                    if (p1.name == "Moon") FindObjectOfType<TrenchInstabilitySimulator>()?.gravitationalForceFromMoon_N += force_on_p2;
                }
            }
        }
    }

    void InitializeDataLogger()
    {
        try
        {
            string filePath = Path.Combine(Application.dataPath, logFileName);
            _dataLogger = new StreamWriter(filePath, false); // false to overwrite existing file
            // Write CSV header
            _dataLogger.WriteLine("Time_s,BodyName,Position_X_m,Position_Y_m,Position_Z_m," +
                                  "Velocity_X_mps,Velocity_Y_mps,Velocity_Z_mps," +
                                  "Acceleration_X_mps2,Acceleration_Y_mps2,Acceleration_Z_mps2");
            Debug.Log($"Data logger '{logFileName}' initialized at {filePath}.");
        }
        catch (Exception e)
        {
            Debug.LogError($"Failed to initialize data logger: {e.Message}");
            enableDataLogging = false; // Disable logging on error
        }
    }

    void LogSimulationData()
    {
        if (_dataLogger == null) return;

        double currentTime = Time.time; // Or use a custom simulation time counter

        foreach (var body in _celestialBodies)
        {
            _dataLogger.WriteLine($"{currentTime:F6}," +
                                  $"{body.name}," +
                                  $"{body.currentPosition_m.x:E6},{body.currentPosition_m.y:E6},{body.currentPosition_m.z:E6}," +
                                  $"{body.currentVelocity_m_per_s.x:E6},{body.currentVelocity_m_per_s.y:E6},{body.currentVelocity_m_per_s.z:E6}," +
                                  $"{body.currentAcceleration_m_per_s_sq.x:E6},{body.currentAcceleration_m_per_s_sq.y:E6},{body.currentAcceleration_m_per_s_sq.z:E6}");
        }
    }

    // Close the data logger file when the script is disabled or application quits
    void OnDisable()
    {
        if (_dataLogger != null)
        {
            _dataLogger.Close();
            Debug.Log($"Data logger '{logFileName}' closed.");
        }
    }

    void OnApplicationQuit()
    {
        // Ensure logger is closed even if application quits unexpectedly
        OnDisable();
    }
}

Planetary Instability Model PIM - Copyright (C) 2025 James Pacha