use std::fs::OpenOptions;
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use std::str::FromStr;

use anyhow::Result;
use clap::{AppSettings, Parser, Subcommand, ValueEnum};
use evanescence_core::geometry::point::IPoint;
use evanescence_core::geometry::region::BoundingRegion;
use evanescence_core::geometry::storage::Soa;
use evanescence_core::numerics::monte_carlo::accept_reject::AcceptReject;
use evanescence_core::numerics::monte_carlo::MonteCarlo;
use evanescence_core::orbital::gaussian_mo::{self, GaussianMo, MoCube};
use evanescence_core::orbital::hybrid::{self, Hybrid, Kind};
use evanescence_core::orbital::{AtomicReal, Qn};
use indicatif::{ProgressBar, ProgressStyle};
use serde::Serialize;

#[derive(Parser, Debug)]
#[clap(author, version, about)]
#[clap(global_setting(AppSettings::DeriveDisplayOrder))]
struct Cli {
    #[clap(subcommand)]
    command: Command,
    #[clap(short = 'n')]
    #[clap(value_parser = clap::value_parser!(u64).range(1..))]
    /// Number of points to sample
    count: u64,
    #[clap(short)]
    /// Path of output json, default stdout
    out: Option<PathBuf>,
    #[clap(long)]
    /// Whether to normalize wavefunction magnitudes to [0, 1]
    normalize_magnitudes: bool,
}

#[derive(Subcommand, Debug)]
enum Command {
    #[clap(allow_negative_numbers = true)]
    /// Sample real hydrogen orbitals
    AtomicReal { n: u32, l: u32, m: i32 },
    /// Sample cube data generated by Gaussian
    GaussianMo {
        /// Path to `.cub` file
        cube: PathBuf,
    },
    /// Sample hybridized hydrogen orbitals
    Hybrid {
        #[clap(arg_enum)]
        kind: HybridKinds,
        #[clap(long, conflicts_with = "nth")]
        /// List the individual linear combinations of the given `kind`
        list: bool,
        #[clap(long, default_value_t = 1)]
        /// Use the n-th linear combination
        nth: usize,
    },
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum HybridKinds {
    SP,
    SP2,
    SP3,
    SP3D,
    SP3D2,
}

impl HybridKinds {
    fn get(&self) -> &'static Kind {
        use hybrid::library;
        #[allow(clippy::explicit_auto_deref)] // False positive.
        match self {
            Self::SP => &*library::SP,
            Self::SP2 => &*library::SP2,
            Self::SP3 => &*library::SP3,
            Self::SP3D => &*library::SP3D,
            Self::SP3D2 => &*library::SP3D2,
        }
    }
}

const SAMPLE_BATCH: usize = 4_096;

fn sample_real<P: IPoint<3>, M: MonteCarlo<3, P, Output = f32>>(
    sampler: &mut M,
    count: usize,
    normalize_magnitudes: bool,
) -> Soa<3, f32> {
    let bar = ProgressBar::new(count as _).with_style(
        ProgressStyle::with_template("[{bar:40}] {pos}/{len}")
            .unwrap()
            .progress_chars("=> "),
    );

    let mut samples = Soa::with_capacity(count);
    while samples.len() < count {
        let batch = usize::min(SAMPLE_BATCH, count - samples.len());
        samples.extend(sampler.take(batch));
        bar.inc(batch as _);
    }

    if normalize_magnitudes {
        let max = samples
            .values()
            .iter()
            .map(|p| p.abs())
            .max_by(f32::total_cmp)
            .unwrap();
        for v in samples.values_mut() {
            *v /= max;
        }
    }

    bar.finish();
    eprintln!("Finished sampling in {:.3}s.", bar.elapsed().as_secs_f32());

    samples
}

fn write_json(out: Option<&PathBuf>, json: impl Serialize) -> Result<()> {
    let out_writer: Box<dyn Write> = if let Some(out) = out {
        Box::new(
            OpenOptions::new()
                .create(true)
                .write(true)
                .truncate(true)
                .open(out)?,
        )
    } else {
        Box::new(std::io::stdout())
    };
    let mut writer = BufWriter::new(out_writer);
    let write_json = if out.is_some() {
        serde_json::to_writer
    } else {
        serde_json::to_writer_pretty
    };
    write_json(&mut writer, &json)?;
    writer.flush()?;
    Ok(())
}

#[derive(Serialize)]
struct AtomicRealJson {
    qn: Qn,
    bounding_sphere_radius: f32,
    x: Vec<f32>,
    y: Vec<f32>,
    z: Vec<f32>,
    psi: Vec<f32>,
}

#[derive(Serialize)]
struct GaussianMoJson {
    atoms: Vec<gaussian_mo::Atom>,
    x: Vec<f32>,
    y: Vec<f32>,
    z: Vec<f32>,
    psi: Vec<f32>,
}

#[derive(Serialize)]
struct HybridJson {
    bounding_sphere_radius: f32,
    x: Vec<f32>,
    y: Vec<f32>,
    z: Vec<f32>,
    psi: Vec<f32>,
}

fn main() -> Result<()> {
    let cli = Cli::parse();

    match cli.command {
        Command::AtomicReal { n, l, m } => {
            let qn = Qn::new(n, l, m)?;
            let atomic_real = AtomicReal::new(qn);
            let bounding_sphere_radius = atomic_real.bounding_region().radius;

            let sampler = &mut AcceptReject::new(atomic_real);
            let ([x, y, z], psi) =
                sample_real(sampler, cli.count as _, cli.normalize_magnitudes).into_components();

            let json = AtomicRealJson {
                qn,
                bounding_sphere_radius,
                x,
                y,
                z,
                psi,
            };
            write_json(cli.out.as_ref(), json)?;
        }
        Command::GaussianMo { cube } => {
            let cube_file = std::fs::read_to_string(cube)?;

            let cube = MoCube::from_str(&cube_file)?;
            let atoms = cube.atoms().clone();
            let gaussian_mo = GaussianMo::new(cube);

            let sampler = &mut AcceptReject::new(gaussian_mo);
            let ([x, y, z], psi) =
                sample_real(sampler, cli.count as _, cli.normalize_magnitudes).into_components();

            let json = GaussianMoJson {
                atoms,
                x,
                y,
                z,
                psi,
            };
            write_json(cli.out.as_ref(), json)?;
        }
        Command::Hybrid { kind, nth, list } => {
            let kind = kind.get();
            if list {
                println!("{kind}-hybridized orbitals:");
                for (idx, combination) in kind.combinations().iter().enumerate() {
                    println!("{}: {combination}", idx + 1);
                }
            } else {
                let hybrid = Hybrid::new(kind.combinations()[nth - 1].clone());
                let bounding_sphere_radius = hybrid.bounding_region().radius;

                let sampler = &mut AcceptReject::new(hybrid);
                let ([x, y, z], psi) =
                    sample_real(sampler, cli.count as _, cli.normalize_magnitudes)
                        .into_components();

                let json = HybridJson {
                    bounding_sphere_radius,
                    x,
                    y,
                    z,
                    psi,
                };
                write_json(cli.out.as_ref(), json)?;
            }
        }
    }

    Ok(())
}
