#!/usr/bin/env python3

import argparse
import pandas as pd
import sys
from pathlib import Path

def create_smiles_file(drug_csv_path, metabolite_file_path, output_path="smiles.smi"):
    """
    Create smiles.smi file by combining drug CSV and metabolite file.
    
    Args:
        drug_csv_path: Path to CSV file with 'SMILES' header
        metabolite_file_path: Path to tab-separated file with SMILES and metabolite IDs
        output_path: Output file path (default: smiles.smi)
    """
    
    # Convert to Path objects for better handling
    drug_csv_path = Path(drug_csv_path)
    metabolite_file_path = Path(metabolite_file_path)
    output_path = Path(output_path)
    
    # Check if input files exist
    if not drug_csv_path.exists():
        raise FileNotFoundError(f"Drug CSV file not found: {drug_csv_path}")
    
    if not metabolite_file_path.exists():
        raise FileNotFoundError(f"Metabolite file not found: {metabolite_file_path}")
    
    # Read drug CSV
    try:
        drug_df = pd.read_csv(drug_csv_path)
    except Exception as e:
        raise Exception(f"Error reading drug CSV file: {e}")
    
    # Validate drug CSV has SMILES column
    if 'SMILES' not in drug_df.columns:
        raise ValueError("Drug CSV file must have 'SMILES' column")
    
    # Check if metabolite file contains "NO_METABOLITES"
    try:
        with open(metabolite_file_path, 'r') as f:
            first_line = f.readline().strip()
    except Exception as e:
        raise Exception(f"Error reading metabolite file: {e}")
    
    # Handle no metabolites case
    if first_line == "NO_METABOLITES":
        print("No metabolites found, creating output file with drugs only")
        
        # Create drug dataframe with drug_N format
        drug_output = drug_df.copy()
        drug_output['ID'] = [f'drug_{i}' for i in range(len(drug_df))]
        drug_output = drug_output[['SMILES', 'ID']]
        
        # Write only drugs to output file
        try:
            drug_output.to_csv(output_path, sep='\t', header=False, index=False)
        except Exception as e:
            raise Exception(f"Error writing output file: {e}")
        
        print(f"Successfully created {output_path}")
        print(f"  - {len(drug_df)} drugs")
        print(f"  - 0 metabolites")
        return
    
    # Read metabolite file normally (tab-separated, no header)
    try:
        metabolite_df = pd.read_csv(metabolite_file_path, sep='\t', header=None, names=['SMILES', 'ID'])
    except Exception as e:
        raise Exception(f"Error reading metabolite file: {e}")
    
    # Validate metabolite file has at least 2 columns
    if metabolite_df.shape[1] < 2:
        raise ValueError("Metabolite file must have at least 2 columns (SMILES and ID)")
    
    # Create drug dataframe with drug_N format
    drug_output = drug_df.copy()
    drug_output['ID'] = [f'drug_{i}' for i in range(len(drug_df))]
    drug_output = drug_output[['SMILES', 'ID']]
    
    # Use metabolite dataframe as-is (already has SMILES and ID columns)
    metabolite_output = metabolite_df[['SMILES', 'ID']]
    
    # Combine drug and metabolite dataframes
    combined_df = pd.concat([drug_output, metabolite_output], ignore_index=True)
    
    # Write to output file (tab-separated, no header)
    try:
        combined_df.to_csv(output_path, sep='\t', header=False, index=False)
    except Exception as e:
        raise Exception(f"Error writing output file: {e}")
    
    print(f"Successfully created {output_path}")
    print(f"  - {len(drug_df)} drugs")
    print(f"  - {len(metabolite_df)} metabolites")

def main():
    parser = argparse.ArgumentParser(
        description="Create smiles.smi file by combining drug CSV and metabolite file",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python3 create_smiles.py drug.csv metabolites.txt
  python3 create_smiles.py drug.csv metabolites.txt -o compounds.smi
        """
    )
    
    parser.add_argument(
        'drug_csv', 
        help='Path to CSV file with SMILES column containing drug compounds'
    )
    
    parser.add_argument(
        'metabolite_file', 
        help='Path to tab-separated file with SMILES and metabolite IDs'
    )
    
    parser.add_argument(
        '-o', '--output', 
        default='smiles.smi',
        help='Output file path (default: smiles.smi)'
    )
    
    args = parser.parse_args()
    
    try:
        create_smiles_file(args.drug_csv, args.metabolite_file, args.output)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()