use charming::{
    element::{Emphasis, EmphasisFocus},
    series::{Sankey, SankeyLink, SankeyNode},
    Chart,
};
use hashbrown::HashMap;
use log::info;
use rayon::prelude::*;
use std::{collections::HashSet, fs::File, sync::Arc};

use crate::{
    annotation::{Annotation, Annotations, AnnotationsStats, Caller},
    callers::{clairs::ClairS, deep_variant::DeepVariant},
    collection::{Initialize, InitializeSolo},
    config::Config,
    io::vcf::write_vcf,
    runners::Run,
    variant::{
        variant::{load_variants, parallel_intersection, RunnerVariants, VcfVariant},
        variant_collection::VariantCollection,
    },
};

pub struct Somatic {
    pub id: String,
    pub config: Config,
    pub annotations: Annotations,
}

impl Initialize for Somatic {
    fn initialize(id: &str, config: crate::config::Config) -> anyhow::Result<Self> {
        let id = id.to_string();
        Ok(Self {
            id,
            config,
            annotations: Annotations::default(),
        })
    }
}

#[derive(Debug, Default, Clone)]
struct SomaticStats {
    initial: HashMap<String, usize>,
    annotations_stats: Vec<AnnotationsStats>,
}

impl SomaticStats {
    pub fn init(collections: &[VariantCollection]) -> Self {
        let mut initial = HashMap::new();

        for collection in collections.iter() {
            let name = format!("{}_{}", collection.vcf.caller, collection.vcf.time);
            initial.insert(name, collection.variants.len());
        }

        Self {
            initial,
            annotations_stats: Default::default(),
        }
    }

    pub fn push_annotations_stats(&mut self, annotations_stats: AnnotationsStats) {
        self.annotations_stats.push(annotations_stats);
    }

    pub fn aggregation(&mut self) -> anyhow::Result<()> {
        let annotations_stats = self
            .annotations_stats
            .get(0)
            .ok_or(anyhow::anyhow!("Can't find stats"))?;
        let step_cat = vec![
            (Annotation::Germline, 0),
            (Annotation::Somatic, 0),
            (Annotation::SoloDiag, 0),
            (Annotation::SoloConstit, 0),
        ];

        let mut callers_cats: Vec<((Annotation, Annotation), Vec<(Annotation, u64)>)> = self
            .initial
            .keys()
            .map(|s| (to_callers_cat(s), step_cat.clone()))
            .collect();

        let mut node_names = Vec::new();

        node_names.extend(step_cat.clone().into_iter().map(|(cat, _)| cat.to_string()));

        let stats: anyhow::Result<Vec<()>> = annotations_stats
            .categorical
            .iter()
            .map(|e| {
                let v = e.value();
                let keys: Vec<&str> = e.key().split(" + ").collect();
                let k_a: Vec<Annotation> = keys
                    .into_iter()
                    .map(|e| e.parse())
                    .collect::<anyhow::Result<_>>()?;

                for ((caller, cat), counts) in callers_cats.iter_mut() {
                    node_names.push(format!("{} {}", caller, cat));
                    if k_a.contains(caller) && k_a.contains(cat) {
                        for (c_annot, value) in counts.iter_mut() {
                            if k_a.contains(c_annot) {
                                *value += v;
                            }
                        }
                    }
                }

                Ok(())
            })
            .collect();
        stats?;

        println!("{callers_cats:#?}");
        let mut links: Vec<(String, String, f64)> = callers_cats
            .iter()
            .flat_map(|((caller, cat), counts)| {
                let from = format!("{} {}", caller, cat);
                counts
                    .iter()
                    .map(move |(annot, count)| (from.clone(), annot.to_string(), *count as f64))
            })
            .collect();
        links.sort_by(|a, b| {
            a.2.partial_cmp(&b.2)
                .unwrap()
                .then(a.0.cmp(&b.0))
                .then(a.1.cmp(&b.1))
        });

        links.dedup();

        node_names.sort();
        node_names.dedup();

        println!("{links:?}");
        println!("{node_names:?}");

        let chart = Chart::new().series(
            Sankey::new()
                .emphasis(Emphasis::new().focus(EmphasisFocus::Adjacency))
                // .data(vec!["a", "b", "a1", "a2", "b1", "c"])
                .data(node_names)
                .links(links), // .links(vec![
                               //     ("a", "a1", 5),
                               //     ("a", "a2", 3),
                               //     ("b", "b1", 8),
                               //     ("a", "b1", 3),
                               //     ("b1", "a1", 1),
                               //     ("b1", "c", 2),
                               // ]),
                               // .data(node_names)
                               // .links(vec![("ClairS Somatic", "Germline", 10)])
                               //     ,
                               // ),
        );

        let mut renderer = charming::ImageRenderer::new(1000, 800);
        renderer.save(&chart, "/data/sankey.svg")?;

        Ok(())
    }
}

pub fn to_callers_cat(s: &str) -> (Annotation, Annotation) {
    let splits: Vec<&str> = s.split("_").collect();

    if splits.contains(&"clairs") {
        (Annotation::Callers(Caller::ClairS), Annotation::Somatic)
    } else if splits.contains(&"clair3-germline") {
        (Annotation::Callers(Caller::ClairS), Annotation::Germline)
    } else if splits.contains(&"DeepVariant") && splits.contains(&"mrd") {
        (
            Annotation::Callers(Caller::DeepVariant),
            Annotation::SoloConstit,
        )
    } else if splits.contains(&"DeepVariant") && splits.contains(&"diag") {
        (
            Annotation::Callers(Caller::DeepVariant),
            Annotation::SoloDiag,
        )
    } else {
        panic!("unknown caller: {s}");
    }
}

impl Run for Somatic {
    fn run(&mut self) -> anyhow::Result<()> {
        let id = self.id.clone();
        info!("Running somatic pipe for {id}.");

        info!("Initialization...");
        let mut v: Vec<Box<dyn RunnerVariants + Send + Sync>> = vec![
            Box::new(ClairS::initialize(&id, self.config.clone())?),
            Box::new(DeepVariant::initialize(&id, "diag", self.config.clone())?),
            Box::new(DeepVariant::initialize(&id, "mrd", self.config.clone())?),
        ];

        let annotations = Arc::new(self.annotations.clone());

        let mut variants_collection = load_variants(&mut v, &annotations)?;
        let clairs_germline =
            ClairS::initialize(&id, self.config.clone())?.germline(&annotations)?;
        variants_collection.push(clairs_germline);
        info!("Variants sources loaded: {}", variants_collection.len());

        let mut somatic_stats = SomaticStats::init(&variants_collection);

        // Annotations Stats
        // let mut annotations = Arc::unwrap_or_clone(annotations);
        somatic_stats.push_annotations_stats(annotations.callers_stat());
        somatic_stats.aggregation()?;

        return Ok(());
        // TODO: look at variants: ClairS + DeepVariant + SoloConstit + SoloDiag + Somatic (error
        // in ClairS somatic)

        // Filtering Somatic variants
        info!("Filtering somatic variants (variants not in salo callers on constit sample or germline).");
        let germline_or_somatic_keys = annotations.get_keys_filter(|anns| {
            !anns.contains(&Annotation::Germline) && !anns.contains(&Annotation::SoloConstit)
        });

        let (somatic_keys, germline_keys, remains) = parallel_intersection(
            &variants_collection
                .iter()
                .flat_map(|e| e.keys())
                .collect::<Vec<u128>>(),
            &germline_or_somatic_keys,
        );
        assert_eq!(0, remains.len());

        info!("Somatic variants positions {}.", somatic_keys.len());
        info!("Germline variants positions {}.", germline_keys.len());

        let somatic_keys: HashSet<u128> = somatic_keys.into_iter().collect();
        annotations.retain_keys(&somatic_keys);
        annotations.callers_stat();

        variants_collection.par_iter_mut().for_each(|c| {
            let before = c.variants.len();
            c.retain_keys(&somatic_keys);
            let after = c.variants.len();
            info!(
                "Variants removed from {}: {}",
                c.vcf.path.display(),
                before - after
            );
        });

        variants_collection.retain(|e| !e.variants.is_empty());

        info!("Constit Bam annotation...");
        variants_collection.iter().try_for_each(|c| {
            c.annotate_with_constit_bam(&annotations, &self.config.normal_bam(&id), 150)
        })?;
        annotations.solo_constit_boundaries(
            self.config.solo_max_alt_constit,
            self.config.solo_min_constit_depth,
        );

        info!("Entropy annotation...");
        variants_collection.iter().for_each(|c| {
            c.annotate_with_sequence_entropy(&annotations, &self.config.reference, 10, 150);
        });
        annotations.callers_stat();

        let prob_keys: HashSet<u128> = annotations
            .get_keys_filter(|anns| {
                // let contains = anns.iter().any(|item| matches!(item, Annotation::SoloDiag));
                // let contains_not = anns.iter().all(|item| !matches!(item, Annotation::Somatic));
                //
                // contains && contains_not
                anns.contains(&Annotation::SoloDiag)
                    && !anns.contains(&Annotation::Somatic)
                    && !anns.contains(&Annotation::LowConstitDepth)
                    && !anns.contains(&Annotation::HighConstitAlt)
            })
            .into_iter()
            .collect();

        info!("Problematic variants {}", prob_keys.len());

        let mut problematic_variants = variants_collection.clone();

        let problematic_variants: Vec<VcfVariant> = problematic_variants
            .iter_mut()
            .flat_map(|e| {
                e.retain_keys(&prob_keys);
                e.variants.clone()
            })
            .collect();

        write_vcf(&problematic_variants, "prob.vcf.gz")?;

        Ok(())
    }
}

// 0-based position
pub fn sequence_at(
    fasta_reader: &mut noodles_fasta::IndexedReader<noodles_fasta::io::BufReader<File>>,
    contig: &str,
    position: usize,
    len: usize,
) -> anyhow::Result<String> {
    // convert to 1-based
    let position = position + 1;

    let start = position.saturating_sub(len / 2).max(1);
    let end = start + len - 1;
    // debug!("Region {contig}:{start}-{end} (1-based inclusive)");

    let start = noodles_core::Position::try_from(start)?;
    let end = noodles_core::Position::try_from(end)?;
    let interval = noodles_core::region::interval::Interval::from(start..=end);

    let r = noodles_core::Region::new(contig.to_string(), interval);
    let record = fasta_reader.query(&r)?;
    let s = String::from_utf8(record.sequence().as_ref().to_vec())?.to_uppercase();

    Ok(s)
}
