#!/usr/bin/env python3
"""
PyTorch Reference Implementation for Pennes Bioheat PINN

This script provides a reference implementation of the Physics-Informed Neural Network
for solving the Pennes Bioheat Equation, used to benchmark against RustyTorch++.

Pennes Bioheat Equation:
    rho*c * dT/dt = k * laplacian(T) + omega_b*rho_b*c_b*(T_a - T) + Q_m + Q_s

Where:
    - T: Temperature field T(x,y,z,t)
    - rho, c: Tissue density and specific heat
    - k: Thermal conductivity
    - omega_b: Blood perfusion rate
    - T_a: Arterial blood temperature (37C)
    - Q_m: Metabolic heat generation
    - Q_s: External heat source (ablation probe)

Usage:
    python pytorch_bioheat_reference.py [--device cuda|cpu] [--steps 1000]
"""

import argparse
import time
from dataclasses import dataclass
from typing import Tuple, Optional, Dict, Any

import numpy as np
import torch
import torch.nn as nn


# =============================================================================
# Tissue Properties
# =============================================================================

@dataclass
class TissueProperties:
    """Physical properties for biological tissue."""
    name: str
    density: float          # kg/m^3
    specific_heat: float    # J/(kg*K)
    conductivity: float     # W/(m*K)
    perfusion: float        # 1/s (blood perfusion rate)
    metabolic_heat: float   # W/m^3

    @classmethod
    def liver(cls) -> 'TissueProperties':
        return cls(
            name="Liver",
            density=1060.0,
            specific_heat=3600.0,
            conductivity=0.512,
            perfusion=0.0064,
            metabolic_heat=420.0
        )

    @classmethod
    def kidney(cls) -> 'TissueProperties':
        return cls(
            name="Kidney",
            density=1050.0,
            specific_heat=3890.0,
            conductivity=0.544,
            perfusion=0.0083,
            metabolic_heat=380.0
        )

    @classmethod
    def tumor(cls) -> 'TissueProperties':
        return cls(
            name="Tumor",
            density=1040.0,
            specific_heat=3800.0,
            conductivity=0.55,
            perfusion=0.002,  # Lower perfusion = hotter
            metabolic_heat=500.0
        )


@dataclass
class BloodProperties:
    """Physical properties for blood."""
    density: float = 1060.0      # kg/m^3
    specific_heat: float = 3770.0  # J/(kg*K)
    arterial_temp: float = 37.0   # Celsius


# =============================================================================
# PINN Network Architecture
# =============================================================================

class FourierFeatures(nn.Module):
    """Learnable Fourier Feature encoding for positional encoding."""

    def __init__(self, in_features: int, num_features: int, scale: float = 1.0):
        super().__init__()
        self.num_features = num_features
        # Initialize frequency matrix B ~ N(0, scale^2)
        self.B = nn.Parameter(torch.randn(in_features, num_features) * scale)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input coordinates [batch, in_features]
        Returns:
            Fourier features [batch, 2*num_features] (cos and sin)
        """
        # x @ B: [batch, num_features]
        proj = 2 * np.pi * (x @ self.B)
        return torch.cat([torch.cos(proj), torch.sin(proj)], dim=-1)


class BioheatPINN(nn.Module):
    """
    Physics-Informed Neural Network for Pennes Bioheat Equation.

    Network architecture:
        Input: (x, y, z, t) -> Fourier Features -> MLP -> Temperature T
    """

    def __init__(
        self,
        fourier_features: int = 64,
        fourier_scale: float = 2.0,
        hidden_layers: list = None,
        activation: str = 'tanh'
    ):
        super().__init__()

        if hidden_layers is None:
            hidden_layers = [128, 128, 128, 64]

        # Fourier feature encoding for 4D input (x, y, z, t)
        self.fourier = FourierFeatures(4, fourier_features, fourier_scale)

        # MLP layers
        layers = []
        in_dim = 2 * fourier_features  # cos + sin features

        for hidden_dim in hidden_layers:
            layers.append(nn.Linear(in_dim, hidden_dim))
            if activation == 'tanh':
                layers.append(nn.Tanh())
            elif activation == 'swish':
                layers.append(nn.SiLU())
            elif activation == 'gelu':
                layers.append(nn.GELU())
            in_dim = hidden_dim

        # Output layer (temperature, no activation)
        layers.append(nn.Linear(in_dim, 1))

        self.mlp = nn.Sequential(*layers)

        # Initialize weights using Xavier
        self._init_weights()

    def _init_weights(self):
        for m in self.mlp.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor, y: torch.Tensor,
                z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Forward pass: predict temperature at given coordinates.

        Args:
            x, y, z: Spatial coordinates [batch]
            t: Time coordinate [batch]

        Returns:
            Temperature prediction [batch, 1]
        """
        coords = torch.stack([x, y, z, t], dim=-1)
        features = self.fourier(coords)
        return self.mlp(features)

    def predict(self, coords: torch.Tensor) -> torch.Tensor:
        """Predict temperature from stacked coordinates [batch, 4]."""
        features = self.fourier(coords)
        return self.mlp(features)


# =============================================================================
# Physics Loss Computation
# =============================================================================

class PennesBioheatLoss:
    """
    Computes physics-informed loss for Pennes Bioheat Equation.

    Loss = L_physics + lambda_bc * L_boundary + lambda_ic * L_initial
    """

    def __init__(
        self,
        tissue: TissueProperties,
        blood: BloodProperties = None,
        probe_power: float = 20.0,
        probe_position: Tuple[float, float, float] = (0.0, 0.0, 0.0),
        probe_radius: float = 0.005,
        domain_bounds: Tuple[float, float, float] = (0.1, 0.1, 0.1),
        lambda_bc: float = 1.0,
        lambda_ic: float = 1.0,
    ):
        self.tissue = tissue
        self.blood = blood or BloodProperties()
        self.probe_power = probe_power
        self.probe_position = torch.tensor(probe_position)
        self.probe_radius = probe_radius
        self.domain_bounds = domain_bounds
        self.lambda_bc = lambda_bc
        self.lambda_ic = lambda_ic

        # Precompute constants
        self.rho_c = tissue.density * tissue.specific_heat
        self.k = tissue.conductivity
        self.perfusion_coeff = (
            tissue.perfusion * self.blood.density * self.blood.specific_heat
        )
        self.T_arterial = self.blood.arterial_temp
        self.Q_m = tissue.metabolic_heat

    def heat_source(self, x: torch.Tensor, y: torch.Tensor,
                    z: torch.Tensor) -> torch.Tensor:
        """Gaussian heat source from ablation probe."""
        dx = x - self.probe_position[0]
        dy = y - self.probe_position[1]
        dz = z - self.probe_position[2]
        r_sq = dx**2 + dy**2 + dz**2

        # Gaussian distribution
        sigma = self.probe_radius
        amplitude = self.probe_power / ((2 * np.pi * sigma**2) ** 1.5)
        return amplitude * torch.exp(-r_sq / (2 * sigma**2))

    def physics_residual(
        self,
        model: BioheatPINN,
        x: torch.Tensor,
        y: torch.Tensor,
        z: torch.Tensor,
        t: torch.Tensor,
    ) -> torch.Tensor:
        """
        Compute PDE residual: rho*c*dT/dt - k*laplacian(T) - perfusion - Q_m - Q_s

        Uses automatic differentiation to compute derivatives.
        """
        # Enable gradients for coordinates
        x = x.requires_grad_(True)
        y = y.requires_grad_(True)
        z = z.requires_grad_(True)
        t = t.requires_grad_(True)

        # Forward pass
        T = model(x, y, z, t)

        # First derivatives
        grad_outputs = torch.ones_like(T)

        dT_dx = torch.autograd.grad(T, x, grad_outputs, create_graph=True)[0]
        dT_dy = torch.autograd.grad(T, y, grad_outputs, create_graph=True)[0]
        dT_dz = torch.autograd.grad(T, z, grad_outputs, create_graph=True)[0]
        dT_dt = torch.autograd.grad(T, t, grad_outputs, create_graph=True)[0]

        # Second derivatives (Laplacian)
        d2T_dx2 = torch.autograd.grad(dT_dx, x, grad_outputs, create_graph=True)[0]
        d2T_dy2 = torch.autograd.grad(dT_dy, y, grad_outputs, create_graph=True)[0]
        d2T_dz2 = torch.autograd.grad(dT_dz, z, grad_outputs, create_graph=True)[0]

        laplacian_T = d2T_dx2 + d2T_dy2 + d2T_dz2

        # Heat source
        Q_s = self.heat_source(x, y, z)

        # PDE residual
        # rho*c*dT/dt = k*laplacian(T) + perfusion*(T_a - T) + Q_m + Q_s
        time_term = self.rho_c * dT_dt
        diffusion_term = self.k * laplacian_T
        perfusion_term = self.perfusion_coeff * (self.T_arterial - T.squeeze())
        source_term = self.Q_m + Q_s

        residual = time_term - diffusion_term - perfusion_term - source_term
        return residual

    def boundary_loss(
        self,
        model: BioheatPINN,
        x: torch.Tensor,
        y: torch.Tensor,
        z: torch.Tensor,
        t: torch.Tensor,
    ) -> torch.Tensor:
        """Dirichlet boundary condition: T = T_body at boundaries."""
        T = model(x, y, z, t)
        T_body = self.T_arterial  # Body temperature
        return torch.mean((T.squeeze() - T_body) ** 2)

    def initial_loss(
        self,
        model: BioheatPINN,
        x: torch.Tensor,
        y: torch.Tensor,
        z: torch.Tensor,
    ) -> torch.Tensor:
        """Initial condition: T(x,y,z,0) = T_body everywhere."""
        t = torch.zeros_like(x)
        T = model(x, y, z, t)
        T_body = self.T_arterial
        return torch.mean((T.squeeze() - T_body) ** 2)

    def compute_loss(
        self,
        model: BioheatPINN,
        collocation_points: Dict[str, torch.Tensor],
        boundary_points: Dict[str, torch.Tensor],
        initial_points: Dict[str, torch.Tensor],
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """
        Compute total PINN loss.

        Returns:
            total_loss: Combined loss tensor
            loss_dict: Dictionary with individual loss components
        """
        # Physics loss at collocation points
        residual = self.physics_residual(
            model,
            collocation_points['x'],
            collocation_points['y'],
            collocation_points['z'],
            collocation_points['t'],
        )
        physics_loss = torch.mean(residual ** 2)

        # Boundary loss
        bc_loss = self.boundary_loss(
            model,
            boundary_points['x'],
            boundary_points['y'],
            boundary_points['z'],
            boundary_points['t'],
        )

        # Initial condition loss
        ic_loss = self.initial_loss(
            model,
            initial_points['x'],
            initial_points['y'],
            initial_points['z'],
        )

        # Total loss
        total_loss = physics_loss + self.lambda_bc * bc_loss + self.lambda_ic * ic_loss

        loss_dict = {
            'total': total_loss.item(),
            'physics': physics_loss.item(),
            'boundary': bc_loss.item(),
            'initial': ic_loss.item(),
        }

        return total_loss, loss_dict


# =============================================================================
# Training Loop
# =============================================================================

class BioheatTrainer:
    """Trainer for Pennes Bioheat PINN."""

    def __init__(
        self,
        model: BioheatPINN,
        loss_fn: PennesBioheatLoss,
        device: torch.device,
        learning_rate: float = 1e-3,
        num_collocation: int = 4096,
        num_boundary: int = 1024,
        num_initial: int = 1024,
        domain_bounds: Tuple[float, float, float] = (0.1, 0.1, 0.1),
        t_end: float = 600.0,
    ):
        self.model = model.to(device)
        self.loss_fn = loss_fn
        self.device = device
        self.optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

        self.num_collocation = num_collocation
        self.num_boundary = num_boundary
        self.num_initial = num_initial
        self.domain_bounds = domain_bounds
        self.t_end = t_end

        self.loss_history = []
        self.step_count = 0

    def sample_collocation_points(self) -> Dict[str, torch.Tensor]:
        """Sample random collocation points in domain."""
        lx, ly, lz = self.domain_bounds

        x = (torch.rand(self.num_collocation) - 0.5) * lx
        y = (torch.rand(self.num_collocation) - 0.5) * ly
        z = (torch.rand(self.num_collocation) - 0.5) * lz
        t = torch.rand(self.num_collocation) * self.t_end

        return {
            'x': x.to(self.device),
            'y': y.to(self.device),
            'z': z.to(self.device),
            't': t.to(self.device),
        }

    def sample_boundary_points(self) -> Dict[str, torch.Tensor]:
        """Sample points on domain boundary."""
        lx, ly, lz = self.domain_bounds
        n = self.num_boundary // 6  # 6 faces

        points_x, points_y, points_z = [], [], []

        # Sample from each face
        for face_val, coord_idx in [
            (-lx/2, 0), (lx/2, 0),  # x faces
            (-ly/2, 1), (ly/2, 1),  # y faces
            (-lz/2, 2), (lz/2, 2),  # z faces
        ]:
            if coord_idx == 0:
                x = torch.full((n,), face_val)
                y = (torch.rand(n) - 0.5) * ly
                z = (torch.rand(n) - 0.5) * lz
            elif coord_idx == 1:
                x = (torch.rand(n) - 0.5) * lx
                y = torch.full((n,), face_val)
                z = (torch.rand(n) - 0.5) * lz
            else:
                x = (torch.rand(n) - 0.5) * lx
                y = (torch.rand(n) - 0.5) * ly
                z = torch.full((n,), face_val)

            points_x.append(x)
            points_y.append(y)
            points_z.append(z)

        x = torch.cat(points_x)
        y = torch.cat(points_y)
        z = torch.cat(points_z)
        t = torch.rand(len(x)) * self.t_end

        return {
            'x': x.to(self.device),
            'y': y.to(self.device),
            'z': z.to(self.device),
            't': t.to(self.device),
        }

    def sample_initial_points(self) -> Dict[str, torch.Tensor]:
        """Sample points at t=0 for initial condition."""
        lx, ly, lz = self.domain_bounds

        x = (torch.rand(self.num_initial) - 0.5) * lx
        y = (torch.rand(self.num_initial) - 0.5) * ly
        z = (torch.rand(self.num_initial) - 0.5) * lz

        return {
            'x': x.to(self.device),
            'y': y.to(self.device),
            'z': z.to(self.device),
        }

    def train_step(self) -> Dict[str, float]:
        """Execute one training step."""
        self.model.train()
        self.optimizer.zero_grad()

        # Sample points
        collocation = self.sample_collocation_points()
        boundary = self.sample_boundary_points()
        initial = self.sample_initial_points()

        # Compute loss
        loss, loss_dict = self.loss_fn.compute_loss(
            self.model, collocation, boundary, initial
        )

        # Backprop
        loss.backward()
        self.optimizer.step()

        self.step_count += 1
        loss_dict['step'] = self.step_count
        self.loss_history.append(loss_dict)

        return loss_dict

    def train(self, num_steps: int, verbose: bool = True) -> list:
        """Train for multiple steps."""
        losses = []

        for i in range(num_steps):
            loss_dict = self.train_step()
            losses.append(loss_dict)

            if verbose and (i + 1) % 100 == 0:
                print(f"Step {self.step_count}: "
                      f"total={loss_dict['total']:.4e}, "
                      f"physics={loss_dict['physics']:.4e}, "
                      f"boundary={loss_dict['boundary']:.4e}")

        return losses


# =============================================================================
# Benchmarking
# =============================================================================

def benchmark_pytorch(
    device: str = 'cuda',
    num_steps: int = 1000,
    fourier_features: int = 64,
    hidden_layers: list = None,
    num_collocation: int = 4096,
) -> Dict[str, Any]:
    """
    Run benchmark for PyTorch Bioheat PINN.

    Returns:
        Dictionary with timing and performance metrics
    """
    if hidden_layers is None:
        hidden_layers = [128, 128, 128, 64]

    # Setup device
    if device == 'cuda' and not torch.cuda.is_available():
        print("CUDA not available, falling back to CPU")
        device = 'cpu'

    device = torch.device(device)
    print(f"Running benchmark on: {device}")

    # Create model
    model = BioheatPINN(
        fourier_features=fourier_features,
        hidden_layers=hidden_layers,
    )

    # Count parameters
    num_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {num_params:,}")

    # Create loss function
    tissue = TissueProperties.liver()
    loss_fn = PennesBioheatLoss(
        tissue=tissue,
        probe_power=20.0,
        probe_position=(0.0, 0.0, 0.0),
    )

    # Create trainer
    trainer = BioheatTrainer(
        model=model,
        loss_fn=loss_fn,
        device=device,
        num_collocation=num_collocation,
    )

    # Warmup
    print("Warming up...")
    for _ in range(10):
        trainer.train_step()

    if device.type == 'cuda':
        torch.cuda.synchronize()

    # Benchmark training
    print(f"Benchmarking {num_steps} training steps...")
    start_time = time.perf_counter()

    trainer.train(num_steps, verbose=True)

    if device.type == 'cuda':
        torch.cuda.synchronize()

    training_time = time.perf_counter() - start_time
    steps_per_second = num_steps / training_time

    # Benchmark inference
    print("Benchmarking inference...")
    model.eval()

    # Create inference grid (32^3 = 32768 points)
    grid_size = 32
    x = torch.linspace(-0.05, 0.05, grid_size)
    y = torch.linspace(-0.05, 0.05, grid_size)
    z = torch.linspace(-0.05, 0.05, grid_size)
    xx, yy, zz = torch.meshgrid(x, y, z, indexing='ij')

    coords = torch.stack([
        xx.flatten(),
        yy.flatten(),
        zz.flatten(),
        torch.full((grid_size**3,), 300.0),  # t = 300s
    ], dim=-1).to(device)

    # Warmup inference
    with torch.no_grad():
        for _ in range(5):
            _ = model.predict(coords)

    if device.type == 'cuda':
        torch.cuda.synchronize()

    # Time inference
    num_inference_runs = 100
    start_time = time.perf_counter()

    with torch.no_grad():
        for _ in range(num_inference_runs):
            _ = model.predict(coords)

    if device.type == 'cuda':
        torch.cuda.synchronize()

    inference_time = time.perf_counter() - start_time
    inference_ms = (inference_time / num_inference_runs) * 1000

    # Memory usage
    if device.type == 'cuda':
        peak_memory_mb = torch.cuda.max_memory_allocated() / (1024 ** 2)
    else:
        peak_memory_mb = 0.0

    # Final loss
    final_loss = trainer.loss_history[-1]['total']

    results = {
        'device': str(device),
        'num_parameters': num_params,
        'num_steps': num_steps,
        'training_time_sec': training_time,
        'steps_per_second': steps_per_second,
        'inference_grid_size': grid_size ** 3,
        'inference_mean_ms': inference_ms,
        'peak_memory_mb': peak_memory_mb,
        'final_loss': final_loss,
    }

    return results


def print_benchmark_results(results: Dict[str, Any]):
    """Pretty print benchmark results."""
    print("\n" + "=" * 60)
    print("PyTorch Bioheat PINN Benchmark Results")
    print("=" * 60)
    print(f"Device:              {results['device']}")
    print(f"Parameters:          {results['num_parameters']:,}")
    print(f"Training steps:      {results['num_steps']}")
    print("-" * 60)
    print(f"Training time:       {results['training_time_sec']:.2f} sec")
    print(f"Steps/second:        {results['steps_per_second']:.1f}")
    print(f"Inference (32^3):    {results['inference_mean_ms']:.2f} ms")
    if results['peak_memory_mb'] > 0:
        print(f"Peak GPU memory:     {results['peak_memory_mb']:.1f} MB")
    print(f"Final loss:          {results['final_loss']:.4e}")
    print("=" * 60)


# =============================================================================
# Main
# =============================================================================

def main():
    parser = argparse.ArgumentParser(
        description='PyTorch Reference Implementation for Pennes Bioheat PINN'
    )
    parser.add_argument(
        '--device', type=str, default='cuda',
        choices=['cuda', 'cpu'],
        help='Device to run on (default: cuda)'
    )
    parser.add_argument(
        '--steps', type=int, default=1000,
        help='Number of training steps (default: 1000)'
    )
    parser.add_argument(
        '--fourier-features', type=int, default=64,
        help='Number of Fourier features (default: 64)'
    )
    parser.add_argument(
        '--collocation', type=int, default=4096,
        help='Number of collocation points (default: 4096)'
    )
    parser.add_argument(
        '--output', type=str, default=None,
        help='Output JSON file for results'
    )

    args = parser.parse_args()

    results = benchmark_pytorch(
        device=args.device,
        num_steps=args.steps,
        fourier_features=args.fourier_features,
        num_collocation=args.collocation,
    )

    print_benchmark_results(results)

    if args.output:
        import json
        with open(args.output, 'w') as f:
            json.dump(results, f, indent=2)
        print(f"\nResults saved to: {args.output}")


if __name__ == '__main__':
    main()