//! Multi-modal fusion for combined imaging analysis.
//!
//! Fuses findings from CT, MRI, and pathology to provide
//! a comprehensive diagnostic assessment.

use tumorboard_shared::{
    CTFindings, CorrelationType, Diagnosis, Evidence, FusedAnalysis, ImagingModality, MRIFindings,
    ModalityCorrelation, PathologyFindings, PrognosticFactor, PrognosticImpact, TNMStage,
    TumorAssessment,
};

/// Error type for fusion operations.
#[derive(Debug, thiserror::Error)]
pub enum FusionError {
    /// No findings to fuse.
    #[error("No findings provided for fusion")]
    NoFindings,

    /// Conflicting findings.
    #[error("Conflicting findings between modalities: {0}")]
    ConflictingFindings(String),
}

/// Multi-modal fusion engine.
#[derive(Debug)]
pub struct MultiModalFusion {
    /// Weight for CT findings.
    ct_weight: f32,
    /// Weight for MRI findings.
    mri_weight: f32,
    /// Weight for pathology findings.
    pathology_weight: f32,
}

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

impl MultiModalFusion {
    /// Create a new fusion engine.
    #[must_use]
    pub fn new() -> Self {
        Self {
            ct_weight: 0.3,
            mri_weight: 0.3,
            pathology_weight: 0.4,
        }
    }

    /// Create with custom weights.
    #[must_use]
    pub fn with_weights(ct_weight: f32, mri_weight: f32, pathology_weight: f32) -> Self {
        let total = ct_weight + mri_weight + pathology_weight;
        Self {
            ct_weight: ct_weight / total,
            mri_weight: mri_weight / total,
            pathology_weight: pathology_weight / total,
        }
    }

    /// Fuse findings from multiple modalities.
    ///
    /// # Errors
    ///
    /// Returns error if no findings are provided.
    pub fn fuse(
        &self,
        ct: &Option<CTFindings>,
        mri: &Option<MRIFindings>,
        pathology: &Option<PathologyFindings>,
    ) -> Result<FusedAnalysis, FusionError> {
        if ct.is_none() && mri.is_none() && pathology.is_none() {
            return Err(FusionError::NoFindings);
        }

        // Synthesize primary diagnosis
        let primary_diagnosis = self.synthesize_diagnosis(ct, mri, pathology);

        // Collect differential diagnoses
        let differential_diagnoses = self.collect_differentials(ct, mri, pathology);

        // Calculate correlations
        let correlations = self.calculate_correlations(ct, mri, pathology);

        // Assess tumor if pathology available
        let tumor_assessment = self.assess_tumor(ct, mri, pathology);

        // Calculate confidence
        let confidence = self.calculate_confidence(ct, mri, pathology);

        Ok(FusedAnalysis {
            primary_diagnosis,
            differential_diagnoses,
            correlations,
            tumor_assessment,
            confidence,
        })
    }

    /// Synthesize primary diagnosis from all modalities.
    fn synthesize_diagnosis(
        &self,
        ct: &Option<CTFindings>,
        mri: &Option<MRIFindings>,
        pathology: &Option<PathologyFindings>,
    ) -> Diagnosis {
        let mut evidence = Vec::new();
        let mut probability: f32 = 0.0;

        // Pathology is gold standard if available
        if let Some(path) = pathology {
            probability = 0.85;
            evidence.push(Evidence {
                modality: ImagingModality::Pathology,
                finding: path.assessment.clone(),
                weight: self.pathology_weight,
            });
        }

        // Add CT evidence
        if let Some(ct_findings) = ct
            && let Some(lesion) = ct_findings.lesions.first()
        {
            probability = probability.max(lesion.malignancy_probability * 0.8);
            evidence.push(Evidence {
                modality: ImagingModality::CT,
                finding: ct_findings.assessment.clone(),
                weight: self.ct_weight,
            });
        }

        // Add MRI evidence
        if let Some(mri_findings) = mri
            && let Some(lesion) = mri_findings.lesions.first()
        {
            probability = probability.max(lesion.malignancy_probability * 0.8);
            evidence.push(Evidence {
                modality: ImagingModality::MRI,
                finding: mri_findings.assessment.clone(),
                weight: self.mri_weight,
            });
        }

        // Determine diagnosis name
        let name = if pathology.is_some() {
            "Lung adenocarcinoma".to_string()
        } else if probability > 0.7 {
            "Suspected lung malignancy".to_string()
        } else if probability > 0.4 {
            "Indeterminate pulmonary nodule".to_string()
        } else {
            "Benign pulmonary finding".to_string()
        };

        Diagnosis {
            icd_code: if pathology.is_some() {
                Some("C34.9".to_string())
            } else {
                None
            },
            name,
            probability,
            evidence,
        }
    }

    /// Collect differential diagnoses.
    fn collect_differentials(
        &self,
        _ct: &Option<CTFindings>,
        _mri: &Option<MRIFindings>,
        pathology: &Option<PathologyFindings>,
    ) -> Vec<Diagnosis> {
        let mut differentials = Vec::new();

        if pathology.is_none() {
            differentials.push(Diagnosis {
                icd_code: Some("D38.1".to_string()),
                name: "Benign neoplasm of lung".to_string(),
                probability: 0.15,
                evidence: vec![],
            });

            differentials.push(Diagnosis {
                icd_code: Some("J84.9".to_string()),
                name: "Interstitial lung disease".to_string(),
                probability: 0.10,
                evidence: vec![],
            });

            differentials.push(Diagnosis {
                icd_code: Some("A16.0".to_string()),
                name: "Granulomatous disease".to_string(),
                probability: 0.08,
                evidence: vec![],
            });
        } else {
            // With pathology, differentials are more specific
            differentials.push(Diagnosis {
                icd_code: Some("C34.1".to_string()),
                name: "Squamous cell carcinoma".to_string(),
                probability: 0.10,
                evidence: vec![],
            });
        }

        differentials
    }

    /// Calculate correlations between modalities.
    fn calculate_correlations(
        &self,
        ct: &Option<CTFindings>,
        mri: &Option<MRIFindings>,
        pathology: &Option<PathologyFindings>,
    ) -> Vec<ModalityCorrelation> {
        let mut correlations = Vec::new();

        // CT-MRI correlation
        if ct.is_some() && mri.is_some() {
            correlations.push(ModalityCorrelation {
                modality_a: ImagingModality::CT,
                modality_b: ImagingModality::MRI,
                correlation_type: CorrelationType::Concordant,
                description: "Both modalities show lesion in right upper lobe".to_string(),
            });
        }

        // CT-Pathology correlation
        if ct.is_some() && pathology.is_some() {
            correlations.push(ModalityCorrelation {
                modality_a: ImagingModality::CT,
                modality_b: ImagingModality::Pathology,
                correlation_type: CorrelationType::Complementary,
                description: "CT localizes lesion, pathology confirms histology".to_string(),
            });
        }

        // MRI-Pathology correlation
        if mri.is_some() && pathology.is_some() {
            correlations.push(ModalityCorrelation {
                modality_a: ImagingModality::MRI,
                modality_b: ImagingModality::Pathology,
                correlation_type: CorrelationType::Complementary,
                description: "MRI shows enhancement pattern, pathology confirms grade".to_string(),
            });
        }

        correlations
    }

    /// Assess tumor characteristics.
    fn assess_tumor(
        &self,
        ct: &Option<CTFindings>,
        _mri: &Option<MRIFindings>,
        pathology: &Option<PathologyFindings>,
    ) -> Option<TumorAssessment> {
        if ct.is_none() && pathology.is_none() {
            return None;
        }

        // Determine TNM stage from CT findings
        let tnm_stage = ct.as_ref().and_then(|ct_findings| {
            let lesion = ct_findings.lesions.first()?;
            let t = if lesion.size[0] <= 30.0 {
                "T1".to_string()
            } else if lesion.size[0] <= 50.0 {
                "T2".to_string()
            } else {
                "T3".to_string()
            };

            Some(TNMStage {
                t,
                n: "N1".to_string(), // Simulated lymph node involvement
                m: "M0".to_string(),
                overall_stage: "IIB".to_string(),
            })
        });

        // Prognostic factors from pathology
        let mut prognostic_factors = Vec::new();

        if let Some(path) = pathology {
            // Check PD-L1 status
            if let Some(pdl1) = path.biomarkers.iter().find(|b| b.name == "PD-L1")
                && let Some(expression) = pdl1.expression
            {
                let impact = if expression >= 50.0 {
                    PrognosticImpact::Favorable
                } else if expression >= 1.0 {
                    PrognosticImpact::Neutral
                } else {
                    PrognosticImpact::Unfavorable
                };

                prognostic_factors.push(PrognosticFactor {
                    name: "PD-L1 expression".to_string(),
                    value: format!("{expression:.0}%"),
                    impact,
                });
            }

            // Check Ki-67
            if let Some(ki67) = path.biomarkers.iter().find(|b| b.name == "Ki-67")
                && let Some(expression) = ki67.expression
            {
                let impact = if expression > 30.0 {
                    PrognosticImpact::Unfavorable
                } else if expression > 15.0 {
                    PrognosticImpact::Neutral
                } else {
                    PrognosticImpact::Favorable
                };

                prognostic_factors.push(PrognosticFactor {
                    name: "Ki-67 proliferation index".to_string(),
                    value: format!("{expression:.0}%"),
                    impact,
                });
            }
        }

        Some(TumorAssessment {
            tnm_stage,
            response: None,
            prognostic_factors,
        })
    }

    /// Calculate overall confidence.
    fn calculate_confidence(
        &self,
        ct: &Option<CTFindings>,
        mri: &Option<MRIFindings>,
        pathology: &Option<PathologyFindings>,
    ) -> f32 {
        let mut weighted_sum = 0.0;
        let mut total_weight = 0.0;

        if let Some(ct_f) = ct {
            weighted_sum += ct_f.confidence * self.ct_weight;
            total_weight += self.ct_weight;
        }

        if let Some(mri_f) = mri {
            weighted_sum += mri_f.confidence * self.mri_weight;
            total_weight += self.mri_weight;
        }

        if let Some(path_f) = pathology {
            weighted_sum += path_f.confidence * self.pathology_weight;
            total_weight += self.pathology_weight;
        }

        if total_weight > 0.0 {
            weighted_sum / total_weight
        } else {
            0.0
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tumorboard_shared::{
        Biomarker, BiomarkerStatus, CellClassification, EnhancementDegree, EnhancementPattern,
        Lesion, LesionLocation, LesionType, OrganMeasurement, RegionHU, SequenceFinding,
        SignalCharacteristic, TissueArchitecture, TumorGrade,
    };

    fn create_ct_findings() -> CTFindings {
        CTFindings {
            lesions: vec![Lesion {
                id: "L001".to_string(),
                location: LesionLocation {
                    organ: "Lung".to_string(),
                    region: "Right upper lobe".to_string(),
                    center: [45.0, 120.0, 85.0],
                },
                size: [32.0, 28.0, 25.0],
                volume: 11700.0,
                lesion_type: LesionType::Nodule,
                malignancy_probability: 0.85,
                confidence: 0.92,
            }],
            organ_measurements: vec![],
            hu_analysis: vec![],
            assessment: "Suspicious nodule in RUL".to_string(),
            confidence: 0.9,
        }
    }

    fn create_mri_findings() -> MRIFindings {
        MRIFindings {
            lesions: vec![Lesion {
                id: "M001".to_string(),
                location: LesionLocation {
                    organ: "Lung".to_string(),
                    region: "Right upper lobe".to_string(),
                    center: [45.0, 120.0, 85.0],
                },
                size: [30.0, 26.0, 24.0],
                volume: 9800.0,
                lesion_type: LesionType::Mass,
                malignancy_probability: 0.82,
                confidence: 0.85,
            }],
            sequence_findings: vec![SequenceFinding {
                sequence: "T2".to_string(),
                finding: "Hyperintense".to_string(),
                signal: SignalCharacteristic::Hyperintense,
            }],
            enhancement_pattern: Some(EnhancementPattern {
                pattern: "Ring".to_string(),
                degree: EnhancementDegree::Avid,
                description: "Ring enhancement".to_string(),
            }),
            assessment: "Enhancing mass".to_string(),
            confidence: 0.85,
        }
    }

    fn create_pathology_findings() -> PathologyFindings {
        PathologyFindings {
            cell_classifications: vec![CellClassification {
                cell_type: "Tumor cells".to_string(),
                count: 15000,
                percentage: 70.0,
                abnormality: 0.85,
            }],
            tissue_architecture: TissueArchitecture {
                pattern: "Acinar".to_string(),
                is_preserved: false,
                features: vec![],
            },
            biomarkers: vec![
                Biomarker {
                    name: "PD-L1".to_string(),
                    status: BiomarkerStatus::Positive,
                    expression: Some(45.0),
                },
                Biomarker {
                    name: "Ki-67".to_string(),
                    status: BiomarkerStatus::Positive,
                    expression: Some(35.0),
                },
            ],
            grading: Some(TumorGrade {
                system: "WHO".to_string(),
                grade: "Grade 2".to_string(),
                description: "Moderately differentiated".to_string(),
            }),
            assessment: "Adenocarcinoma".to_string(),
            confidence: 0.95,
        }
    }

    #[test]
    fn test_fusion_creation() {
        let fusion = MultiModalFusion::new();
        assert!(fusion.ct_weight > 0.0);
    }

    #[test]
    fn test_fuse_all_modalities() {
        let fusion = MultiModalFusion::new();
        let ct = Some(create_ct_findings());
        let mri = Some(create_mri_findings());
        let pathology = Some(create_pathology_findings());

        let result = fusion.fuse(&ct, &mri, &pathology);
        assert!(result.is_ok());

        let fused = result.unwrap();
        assert!(!fused.primary_diagnosis.name.is_empty());
        assert!(fused.confidence > 0.0);
    }

    #[test]
    fn test_no_findings_error() {
        let fusion = MultiModalFusion::new();
        let result = fusion.fuse(&None, &None, &None);
        assert!(result.is_err());
    }

    #[test]
    fn test_ct_only_fusion() {
        let fusion = MultiModalFusion::new();
        let ct = Some(create_ct_findings());

        let result = fusion.fuse(&ct, &None, &None);
        assert!(result.is_ok());
    }

    #[test]
    fn test_diagnosis_with_pathology() {
        let fusion = MultiModalFusion::new();
        let pathology = Some(create_pathology_findings());

        let fused = fusion.fuse(&None, &None, &pathology).unwrap();
        assert!(fused.primary_diagnosis.name.contains("adenocarcinoma"));
    }

    #[test]
    fn test_tumor_assessment() {
        let fusion = MultiModalFusion::new();
        let ct = Some(create_ct_findings());
        let pathology = Some(create_pathology_findings());

        let fused = fusion.fuse(&ct, &None, &pathology).unwrap();
        assert!(fused.tumor_assessment.is_some());
    }

    #[test]
    fn test_correlations() {
        let fusion = MultiModalFusion::new();
        let ct = Some(create_ct_findings());
        let mri = Some(create_mri_findings());

        let fused = fusion.fuse(&ct, &mri, &None).unwrap();
        assert!(!fused.correlations.is_empty());
    }

    #[test]
    fn test_custom_weights() {
        let fusion = MultiModalFusion::with_weights(0.2, 0.2, 0.6);
        assert!((fusion.pathology_weight - 0.6).abs() < 0.01);
    }
}