#!/usr/bin/env python3
import argparse
import subprocess
import sys
from pathlib import Path

def create_symdoc_file(filename: str, twist: float, rise: float):
    """
    Generates the symdoc.dat file with the specified helical parameters.
    The formatting is critical for the 'himpose' program.
    """
    # The symdoc.dat file has a specific format where only the last
    # row contains the relevant helical symmetry parameters.
    content = (
        "    1 2   -4.500000     4.68000\n"
        "    2 2   -4.500000     4.68000\n"
        # Format the twist and rise with specific padding and precision.
        f"    3 2   {twist:11.5f}   {rise:10.4f}\n"
    )
    try:
        with open(filename, 'w') as f:
            f.write(content)
        print(f"✅ Successfully created '{filename}' with Twist={twist}, Rise={rise}")
    except IOError as e:
        print(f"❌ Error: Could not write to file {filename}: {e}")
        sys.exit(1)

def create_cp_spider_file(filename: str):
    """
    Generates the cp-spider.spd file needed for a SPIDER command.
    """
    # The cp-spider.spd file contains a simple SPIDER command
    # sequence to copy a file.
    content = (
        "cp\n"
        "tmp3\n"
        "tmp4\n"
        "en d\n"
    )
    try:
        with open(filename, 'w') as f:
            f.write(content)
        print(f"✅ Successfully created '{filename}'")
    except IOError as e:
        print(f"❌ Error: Could not write to file {filename}: {e}")
        sys.exit(1)

def run_command(command: list):
    """
    Executes a shell command and checks for errors.
    """
    command_str = ' '.join(command)
    print(f"\n🚀 Executing: {command_str}")
    try:
        subprocess.run(command, check=True, text=True)
    except FileNotFoundError:
        print(f"❌ Error: Command '{command[0]}' not found.")
        print("Please ensure the required software (EMAN2, SPIDER, etc.) is in your PATH.")
        sys.exit(1)
    except subprocess.CalledProcessError as e:
        print(f"❌ Error: Command failed with exit code {e.returncode}")
        sys.exit(1)

def main():
    """
    Main function to parse arguments and execute the workflow.
    """
    parser = argparse.ArgumentParser(
        description="""
        A Python script to perform helical imposition on a cryo-EM map.
        This script automates a series of steps using EMAN2 and SPIDER.
        """,
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument("map_name", type=str, help="Input cryo-EM map file name (e.g., my_map.mrc).")
    parser.add_argument("helical_rise", type=float, help="Refined helical rise (in Angstroms).")
    parser.add_argument("helical_twist", type=float, help="Refined helical twist (in degrees).")
    parser.add_argument("pixel_size", type=float, help="Pixel size of the map (in Angstroms/pixel).")
    parser.add_argument("max_diameter", type=float, help="Estimated maximum diameter of the filament (in Angstroms).")

    args = parser.parse_args()

    # --- User Guidance ---
    print("*" * 60)
    print("📢 IMPORTANT: This script requires an active SBGrid environment.")
    print("Please run 'sbg' in your terminal before executing this script.")
    print("*" * 60)

    # --- Step 1: Prepare control files ---
    print("--- Preparing control files ---")
    create_symdoc_file("symdoc.dat", args.helical_twist, args.helical_rise)
    create_cp_spider_file("cp-spider.spd")

    # --- Step 2: Define filenames ---
    input_path = Path(args.map_name)
    output_map = f"{input_path.stem}-himpose.mrc"

    # --- Step 3: Execute processing pipeline ---
    print("\n--- Starting processing pipeline ---")

    # (b) EMReady.sh
    cmd_emready = ["EMReady.sh", args.map_name, "tmp1.mrc", "-g", "0", "-b", "20", "-s", "48"]
    run_command(cmd_emready)

    # (c) e2proc3d.py to convert to SPIDER format
    cmd_e2p_1 = ["e2proc3d.py", "tmp1.mrc", "tmp2.spi", "--outtype=spidersingle"]
    run_command(cmd_e2p_1)

    # (d) e2proc3d.py to rotate the volume
    cmd_e2p_2 = ["e2proc3d.py", "tmp2.spi", "tmp3.spi", "--rot=0,90,0", "--outtype=spidersingle"]
    run_command(cmd_e2p_2)

    # (e) SPIDER copy command
    cmd_spider = ["spider", "spd/spi", "@cp-spider"]
    run_command(cmd_spider)

    # (f) himpose to apply helical symmetry
    cmd_himpose = [
        "himpose", "tmp4.spi", "symdoc.dat", "tmp5.spi",
        f"{args.pixel_size:.2f}", "0.000", f"{args.max_diameter:.4f}"
    ]
    run_command(cmd_himpose)

    # (g) e2proc3d.py to convert back to MRC
    cmd_e2p_3 = ["e2proc3d.py", "tmp5.spi", output_map, f"--apix={args.pixel_size}"]
    run_command(cmd_e2p_3)

    print("\n" + "="*60)
    print(f"🎉 Workflow finished successfully!")
    print(f"Final output map: {output_map}")
    print("Intermediate files (tmp*.mrc, tmp*.spi) were preserved for inspection.")
    print("="*60)


if __name__ == "__main__":
    main()