# We recommend changing the name of this file from "Supplementary_File_1.txt" to "classify_rna.pl".


#!/usr/bin/perl
use strict;
use warnings;
use Getopt::Long;

# Argument Definitions
my ($fasta, $rfam_file, $mirna_bwt, $ncrna_bwt, $outpre);
GetOptions(
    'fa=s'       => \$fasta,
    'rfam=s'     => \$rfam_file,
    'mirna=s'    => \$mirna_bwt,
    'ncrna=s'    => \$ncrna_bwt,
    'outpre=s'   => \$outpre,
) or die "Usage: perl classify_rna.pl -fa <fasta> -rfam <rfam_file> -mirna <miRNA.bwt> -ncrna <ncRNA.bwt> -outpre <output_prefix>\n";

# Hash for classifying lead IDs by priority
my %category;

# Sub: Retrieve read IDs from BWT files and register them in categories (if not already classified)
sub read_bwt_simple {
    my ($file, $label) = @_;
    return unless $file;
    open my $fh, "<", $file or die "Cannot open $file: $!";
    while (<$fh>) {
        next if /^@/;
        my ($id) = split(/\t/);
        $category{$id} = $label unless exists $category{$id};
    }
    close $fh;
}

# Rfam ID → Category Mapping
my %rfam_map;
open my $rf, "<", $rfam_file or die "Cannot open Rfam file: $!";
while (<$rf>) {
    chomp;
    my ($rfam_id, $name, $desc) = split(/\t/);
    if ($desc =~ /rRNA/)      { $rfam_map{$rfam_id} = "rRNA"; }
    elsif ($desc =~ /snRNA/)  { $rfam_map{$rfam_id} = "snRNA"; }
    elsif ($desc =~ /snoRNA/) { $rfam_map{$rfam_id} = "snoRNA"; }
    elsif ($desc =~ /tRNA/)   { $rfam_map{$rfam_id} = "tRNA"; }
    else                      { $rfam_map{$rfam_id} = "other_ncRNA"; }
}
close $rf;

# miRNA is directly categorized
read_bwt_simple($mirna_bwt, "miRNA");

# Rfam-mapped read ID → RfamID mapping
my %read_to_rfam;
if ($ncrna_bwt) {
    open my $fh, "<", $ncrna_bwt or die "Cannot open $ncrna_bwt: $!";
    while (<$fh>) {
        next if /^@/;
        my @cols = split(/\t/);
        my ($read_id, $ref_id) = @cols[0,2];  # Reference name including read ID and Rfam
        my ($rfid) = $ref_id =~ /(RF\d{5})/;
        if ($rfid) {
            $category{$read_id} = "ncRNA" unless exists $category{$read_id};
            $read_to_rfam{$read_id} = $rfid;
        }
    }
    close $fh;
}

# Only target categories for output (those to be converted to FASTA format)
my %valid_categories = map { $_ => 1 } qw(miRNA rRNA snRNA snoRNA tRNA other_ncRNA unknown);

# Classified FASTA File Handles
my %fasta_out;

# Sub: Returns an output file handle corresponding to the classification name (only what is necessary)
sub get_fasta_handle {
    my ($label) = @_;
    return unless exists $valid_categories{$label};
    unless (exists $fasta_out{$label}) {
        my $file = $outpre . "." . $label . ".fasta";
        open my $fh, ">", $file or die "Cannot write $file: $!";
        $fasta_out{$label} = $fh;
    }
    return $fasta_out{$label};
}

# Output file
my $out_file = $outpre . ".classification.txt";
open my $in_fa, "<", $fasta or die "Cannot open $fasta: $!";
open my $out,   ">", $out_file or die "Cannot write $out_file: $!";

# Processing FASTA files while classifying and outputting
my ($id, $seq) = ("", "");
while (<$in_fa>) {
    chomp;
    if (/^>(\S+)/) {
        if ($id ne "") {
            my $label = $category{$id} || "unknown";
            if ($label eq "ncRNA") {
                my $rfid = $read_to_rfam{$id};
                $label = $rfam_map{$rfid} || "other_ncRNA";
            }
            print $out "$id\t$seq\t$label\n";
            if (my $fh = get_fasta_handle($label)) {
                print $fh ">$id\n$seq\n";
            }
        }
        $id = $1;
        $seq = "";
    } else {
        $seq .= $_;
    }
}

# Process the last array
if ($id ne "") {
    my $label = $category{$id} || "unknown";
    if ($label eq "ncRNA") {
        my $rfid = $read_to_rfam{$id};
        $label = $rfam_map{$rfid} || "other_ncRNA";
    }
    print $out "$id\t$seq\t$label\n";
    if (my $fh = get_fasta_handle($label)) {
        print $fh ">$id\n$seq\n";
    }
}

# Close
close $in_fa;
close $out;
foreach my $fh (values %fasta_out) {
    close $fh;
}

print “✔ Output complete: $out_file + FASTA files for each category\n”;