//! Multi-physics coupling for PINNs

use crate::error::Result;
use crate::physics::PhysicsLoss;
use async_trait::async_trait;

/// Multi-physics PINN for coupled physics problems
pub struct MultiPhysicsPINN {
    /// Individual physics models
    pub physics_models: Vec<Box<dyn PhysicsLoss + Send + Sync>>,
    /// Coupling terms
    pub couplings: Vec<Box<dyn PhysicsCoupling + Send + Sync>>,
}

/// Physics coupling trait
#[async_trait]
pub trait PhysicsCoupling: Send + Sync {
    /// Compute coupling term
    async fn compute_coupling(&self) -> Result<f64>;
}

/// Thermal coupling (Boussinesq approximation)
pub struct ThermalCoupling {
    pub gravity: f64,
    pub thermal_expansion: f64,
}

impl ThermalCoupling {
    #[must_use]
    pub fn boussinesq(gravity: f64, thermal_expansion: f64) -> Self {
        Self {
            gravity,
            thermal_expansion,
        }
    }
}

#[async_trait]
impl PhysicsCoupling for ThermalCoupling {
    async fn compute_coupling(&self) -> Result<f64> {
        Ok(self.gravity * self.thermal_expansion)
    }
}

impl MultiPhysicsPINN {
    #[must_use]
    pub fn builder() -> MultiPhysicsBuilder {
        MultiPhysicsBuilder::new()
    }
}

/// Builder for multi-physics PINNs
pub struct MultiPhysicsBuilder {
    physics_models: Vec<Box<dyn PhysicsLoss + Send + Sync>>,
    couplings: Vec<Box<dyn PhysicsCoupling + Send + Sync>>,
}

impl Default for MultiPhysicsBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl MultiPhysicsBuilder {
    #[must_use]
    pub fn new() -> Self {
        Self {
            physics_models: Vec::new(),
            couplings: Vec::new(),
        }
    }

    #[must_use]
    pub fn add_physics(mut self, physics: Box<dyn PhysicsLoss + Send + Sync>) -> Self {
        self.physics_models.push(physics);
        self
    }

    #[must_use]
    pub fn add_coupling(mut self, coupling: Box<dyn PhysicsCoupling + Send + Sync>) -> Self {
        self.couplings.push(coupling);
        self
    }

    pub fn build(self) -> Result<MultiPhysicsPINN> {
        Ok(MultiPhysicsPINN {
            physics_models: self.physics_models,
            couplings: self.couplings,
        })
    }
}