molx 0.1.9

use std::fs::File;
use std::io::{self, BufWriter, IsTerminal};
use std::path::PathBuf;

use clap::{Args, Parser, Subcommand, ValueEnum};
use molx::atlas::{AtlasClient, AtlasConfidence, read_sequence_file, validate_atlas_id, write_pdb};
use molx::{ColorScheme, Protein, Representation, Viewer};
use rstui::{Runner, RunnerConfig};

#[derive(Parser)]
#[command(
    name = "molx",
    version,
    about = "Explore protein structures in a graphical terminal",
    long_about = None
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Open a local PDB structure in the interactive 3D viewer.
    View(ViewArgs),
    /// Print a compact summary of a PDB structure.
    Info {
        /// Local PDB file.
        path: PathBuf,
    },
    /// Render a PDB structure to a PNG without entering the interactive viewer.
    Snapshot(SnapshotArgs),
    /// Fetch predicted structures or run a small ESMFold job.
    Atlas {
        #[command(subcommand)]
        command: AtlasCommand,
    },
}

#[derive(Args)]
struct ViewArgs {
    /// Local PDB file.
    path: PathBuf,
    /// Initial molecular representation.
    #[arg(long, value_enum, default_value_t = RepresentationArg::Ribbon)]
    representation: RepresentationArg,
    /// Initial color scheme; auto selects pLDDT for ESMFold structures.
    #[arg(long, value_enum, default_value_t = ColorArg::Auto)]
    color: ColorArg,
}

#[derive(Args)]
struct SnapshotArgs {
    /// Local PDB file.
    path: PathBuf,
    /// Destination PNG file.
    #[arg(short, long, default_value = "molx.png")]
    output: PathBuf,
    /// Image width in pixels.
    #[arg(long, default_value_t = 1200)]
    width: u32,
    /// Image height in pixels.
    #[arg(long, default_value_t = 800)]
    height: u32,
    /// Molecular representation.
    #[arg(long, value_enum, default_value_t = RepresentationArg::Ribbon)]
    representation: RepresentationArg,
    /// Color scheme; auto selects pLDDT for ESMFold structures.
    #[arg(long, value_enum, default_value_t = ColorArg::Auto)]
    color: ColorArg,
    /// Add an ESM Atlas PAE panel for this MGYP identifier.
    #[arg(long)]
    atlas_id: Option<String>,
    /// Preselect a residue as CHAIN:NUMBER, for example A:121.
    #[arg(long, value_name = "CHAIN:NUMBER")]
    select: Option<String>,
    /// Camera yaw in degrees.
    #[arg(long, allow_hyphen_values = true)]
    yaw: Option<f32>,
    /// Camera pitch in degrees.
    #[arg(long, allow_hyphen_values = true)]
    pitch: Option<f32>,
    /// Camera zoom multiplier.
    #[arg(long)]
    zoom: Option<f32>,
    /// Replace an existing PNG file.
    #[arg(long)]
    force: bool,
}

#[derive(Subcommand)]
enum AtlasCommand {
    /// Download a predicted PDB structure by MGnify/ESM Atlas ID.
    Fetch {
        /// Identifier beginning with MGYP, for example MGYP002537940442.
        id: String,
        /// Destination PDB path; defaults to <ID>.pdb.
        #[arg(short, long)]
        output: Option<PathBuf>,
        /// Open the downloaded structure after saving it.
        #[arg(long)]
        view: bool,
        /// Replace an existing output file.
        #[arg(long)]
        force: bool,
    },
    /// Predict a structure from one protein sequence with the public ESMFold API.
    Fold(FoldArgs),
}

#[derive(Args)]
struct FoldArgs {
    /// Amino-acid sequence as a command-line string.
    #[arg(long, required_unless_present = "fasta", conflicts_with = "fasta")]
    sequence: Option<String>,
    /// FASTA or plain sequence file.
    #[arg(
        long,
        required_unless_present = "sequence",
        conflicts_with = "sequence"
    )]
    fasta: Option<PathBuf>,
    /// Destination PDB path.
    #[arg(short, long, default_value = "esmfold.pdb")]
    output: PathBuf,
    /// Open the prediction after saving it.
    #[arg(long)]
    view: bool,
    /// Replace an existing output file.
    #[arg(long)]
    force: bool,
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum RepresentationArg {
    #[value(alias = "backbone")]
    Ribbon,
    Atoms,
    Combined,
}

impl From<RepresentationArg> for Representation {
    fn from(value: RepresentationArg) -> Self {
        match value {
            RepresentationArg::Ribbon => Self::Backbone,
            RepresentationArg::Atoms => Self::Atoms,
            RepresentationArg::Combined => Self::Combined,
        }
    }
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum ColorArg {
    Auto,
    Chain,
    Element,
    Confidence,
}

fn main() {
    if let Err(error) = run(Cli::parse()) {
        eprintln!("molx: {error}");
        std::process::exit(1);
    }
}

fn run(cli: Cli) -> Result<(), String> {
    match cli.command {
        Command::View(arguments) => view_path(
            arguments.path,
            arguments.representation.into(),
            arguments.color,
            None,
        ),
        Command::Info { path } => {
            println!("{}", Protein::from_path(path)?.summary());
            Ok(())
        }
        Command::Snapshot(arguments) => snapshot(arguments),
        Command::Atlas { command } => run_atlas(command),
    }
}

fn snapshot(arguments: SnapshotArgs) -> Result<(), String> {
    if arguments.output.exists() && !arguments.force {
        return Err(format!(
            "{} already exists; pass --force to replace it",
            arguments.output.display()
        ));
    }
    let protein = Protein::from_path(arguments.path)?;
    let color = resolve_color(arguments.color, &protein);
    let mut viewer = Viewer::new(protein, arguments.representation.into(), color);
    if let Some(id) = arguments.atlas_id {
        let id = validate_atlas_id(&id)?;
        eprintln!("Fetching PAE for {id} from ESM Atlas...");
        let confidence = AtlasClient::new()?.fetch_confidence(&id)?;
        viewer = viewer.with_atlas(id, Some(confidence));
    }
    if let Some(selector) = arguments.select {
        let (chain, number) = parse_residue_selector(&selector)?;
        viewer = viewer.with_selected_residue(&chain, number)?;
    }
    if arguments.yaw.is_some() || arguments.pitch.is_some() {
        viewer = viewer.with_camera_angles(
            arguments.yaw.unwrap_or(-41.25),
            arguments.pitch.unwrap_or(21.77),
        );
    }
    if let Some(zoom) = arguments.zoom {
        viewer = viewer.with_zoom(zoom)?;
    }
    let frame = viewer.snapshot(arguments.width, arguments.height);
    let file = File::create(&arguments.output)
        .map_err(|error| format!("cannot create {}: {error}", arguments.output.display()))?;
    let mut encoder = png::Encoder::new(BufWriter::new(file), frame.width, frame.height);
    encoder.set_color(png::ColorType::Rgba);
    encoder.set_depth(png::BitDepth::Eight);
    let mut writer = encoder
        .write_header()
        .map_err(|error| format!("cannot initialize PNG output: {error}"))?;
    writer
        .write_image_data(&frame.pixels)
        .map_err(|error| format!("cannot write {}: {error}", arguments.output.display()))?;
    println!("Saved {}", arguments.output.display());
    Ok(())
}

fn run_atlas(command: AtlasCommand) -> Result<(), String> {
    match command {
        AtlasCommand::Fetch {
            id,
            output,
            view,
            force,
        } => {
            let id = validate_atlas_id(&id)?;
            let output = output.unwrap_or_else(|| PathBuf::from(format!("{id}.pdb")));
            eprintln!("Fetching {id} from ESM Atlas...");
            let client = AtlasClient::new()?;
            let pdb = client.fetch_structure(&id)?;
            write_pdb(&output, &pdb, force)?;
            eprintln!("Saved {}", output.display());
            if view {
                eprintln!("Fetching predicted aligned error for {id}...");
                let confidence = match client.fetch_confidence(&id) {
                    Ok(confidence) => Some(confidence),
                    Err(error) => {
                        eprintln!("molx: warning: {error}");
                        None
                    }
                };
                view_path(
                    output,
                    Representation::Backbone,
                    ColorArg::Confidence,
                    Some((id, confidence)),
                )
            } else {
                Ok(())
            }
        }
        AtlasCommand::Fold(arguments) => {
            let sequence = if let Some(sequence) = arguments.sequence {
                molx::atlas::normalize_sequence(&sequence)?
            } else if let Some(path) = arguments.fasta {
                read_sequence_file(path)?
            } else {
                unreachable!("clap requires one sequence input")
            };
            eprintln!(
                "Folding {} residues with the public ESMFold service...",
                sequence.len()
            );
            let pdb = AtlasClient::new()?.fold_sequence(&sequence)?;
            write_pdb(&arguments.output, &pdb, arguments.force)?;
            eprintln!("Saved {}", arguments.output.display());
            if arguments.view {
                view_path(
                    arguments.output,
                    Representation::Backbone,
                    ColorArg::Confidence,
                    Some(("Live ESMFold".to_owned(), None)),
                )
            } else {
                Ok(())
            }
        }
    }
}

fn view_path(
    path: PathBuf,
    representation: Representation,
    color_argument: ColorArg,
    atlas: Option<(String, Option<AtlasConfidence>)>,
) -> Result<(), String> {
    if !io::stdout().is_terminal() {
        return Err("the interactive viewer needs a terminal on stdout; use `molx info` for non-interactive output".to_owned());
    }
    let protein = Protein::from_path(path)?;
    let color = resolve_color(color_argument, &protein);
    let mut viewer = Viewer::new(protein, representation, color);
    if let Some((label, confidence)) = atlas {
        viewer = viewer.with_atlas(label, confidence);
    }
    Runner::new(RunnerConfig {
        frame_rate: 30,
        idle_refine_ms: 110,
        ..RunnerConfig::default()
    })
    .run(viewer)
}

fn resolve_color(color_argument: ColorArg, protein: &Protein) -> ColorScheme {
    match color_argument {
        ColorArg::Auto if protein.is_prediction => ColorScheme::Confidence,
        ColorArg::Auto | ColorArg::Chain => ColorScheme::Chain,
        ColorArg::Element => ColorScheme::Element,
        ColorArg::Confidence => ColorScheme::Confidence,
    }
}

fn parse_residue_selector(value: &str) -> Result<(String, i32), String> {
    let (chain, number) = value
        .split_once(':')
        .ok_or_else(|| format!("invalid residue selector {value:?}; expected CHAIN:NUMBER"))?;
    if chain.trim().is_empty() {
        return Err(format!(
            "invalid residue selector {value:?}; chain cannot be empty"
        ));
    }
    let number = number
        .trim()
        .parse::<i32>()
        .map_err(|_| format!("invalid residue selector {value:?}; residue number is not valid"))?;
    Ok((chain.trim().to_owned(), number))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_residue_selector() {
        assert_eq!(
            parse_residue_selector("A:121").unwrap(),
            ("A".to_owned(), 121)
        );
        assert!(parse_residue_selector("121").is_err());
    }
}