#!/usr/bin/env perl   

#  -*- perl -*-

use strict;
use warnings FATAL => qw ( all );
use Data::Dumper;
use Formats;
use Errors;

# Ebbe Sloth Andersen, July 2005.

# STEM_COLORS: This program will shade an RNA alignment according to the
# assigned structures for each sequence. Basically, it converts the large-
# small-letter annotation that is used in the col-format to assign r, g and b
# color columns that can be read by other programs. Stem shade regions
# will appear as blocks in an RNA structural alignment.

# >>>>>>>>>>>>>>>>>> RUN PROGRAM <<<<<<<<<<<<<<<<<<<<

my ( $entries, $file, $header );

( $file ) = @ARGV;

( $header, $entries ) = &Formats::read_col ( $file );
&Errors::check_col ( $entries, { "TYPE" => ["RNA","pairingmask"] } );
$entries = &stem_colors ( $entries );
$header .= "; stem_colors.pl was run on this file.";
&Formats::write_col ( $header, $entries );

# >>>>>>>>>>>>>>>>>> SUBROUTINES <<<<<<<<<<<<<<<<<<<<<<<

sub stem_colors
{
    # Ebbe Sloth Andersen, July 2005.

    my ( $entries ) = @_;

    # Returns an array of hashes.

    my ( @rna,
	 $rna,
	 @rnas,
	 $pm,
         @pm,
	 @pms,
	 @bps,
	 $bp,
         $i,
         $n,
         $m,
         %colors,
         @color,
         $color,
         $prev,
         $prev2,
         $length,
         @labels,
         $label,
         @colo,
         $find,
         @colormask,
         $align_bp,
         $w,
         $change,
         );

    # Define colors
    %colors = (
	       "red3"      => [1.00, 0.50, 0.50],  # lighter colors
	       "orange3"   => [1.00, 0.70, 0.50], 
	       "yellow3"   => [1.00, 1.00, 0.50],
	       "green3"    => [0.50, 1.00, 0.50],
	       "cyan3"     => [0.50, 1.00, 1.00],
	       "blue3"     => [0.50, 0.50, 1.00],
	       "grey3"     => [0.70, 0.70, 0.70],
	       "new3"	   => [0.10,0.50,0.90],
	       );

    # array of colors
    @color = ( 
               "red3", "orange3", "yellow3", "green3", "cyan3", "blue3",    # lighter colors
             );

    @rnas = grep { $_->{'TYPE'} eq "RNA" } @{ $entries };
    @pms = grep { $_->{'TYPE'} eq "pairingmask" } @{ $entries };
    @pm = split ( /,/, $pms[0]->{'residue'} );

    $i = 0;
    $n = -1;
    $prev = "-";;
    $length = scalar @pm;
    @labels = ( );
    $find = "false";
    $color = "new3";

    $pms[0]->{'color_r'} = "";
    $pms[0]->{'color_g'} = "";
    $pms[0]->{'color_b'} = "";        

    foreach $pm ( @pm ) {
        
        if ( $pm eq "-" ) {
            $pms[0]->{'color_r'} .= "1.00,";
            $pms[0]->{'color_g'} .= "1.00,";
            $pms[0]->{'color_b'} .= "1.00,";        
        } else {
            if ( $pm[$i] ne $prev ) {
                $m = 0;
                foreach $label ( @labels ) {        # but search for similar stem
                    if ( $pm[$i] eq $label ) {      # and get its color
                        $find = "true";
                        last;
                    } else {
                        $find = "false";
                    }
                    $m++;
                    if ( $m > 5 ) { $m = 0; }
                }
                if ( $find eq "false" ) {           # change color if new stem
                    $n++;
                    push ( @labels, $pm );
                }
            }
            if ( $n > 5 ) { $n = 0; }
            if ( $find eq "true" ) {
                $color = $color[$m];
            } else {
                $color = $color[$n];
            }
            if ( $pm[$i] ne $prev ) {
                push ( @colo, $color );             # register colors
            }
            $pms[0]->{'color_r'} .= "$colors{$color}[0],";
            $pms[0]->{'color_g'} .= "$colors{$color}[1],";
            $pms[0]->{'color_b'} .= "$colors{$color}[2],";
        }
        if ( $pm[$i] ne "-" ) {
            $prev = $pm[$i];
        }
        push ( @colormask, $color );
        $color = "new3";
        $i++;
    }
    foreach $rna ( @rnas ) {

        $rna->{'color_r'} = "";
        $rna->{'color_g'} = "";
        $rna->{'color_b'} = "";        

	@bps = ( );
	@bps = split(/,/, $rna->{'residue'});
        $i = 0;
        $n = 0;
	foreach $bp ( @bps ) {
	    if ( $bp =~ /^(T|A|U|G|C)/ ) {
                $color = $colormask[$i];
                $rna->{'color_r'} .= "$colors{$color}[0],";
                $rna->{'color_g'} .= "$colors{$color}[1],";
                $rna->{'color_b'} .= "$colors{$color}[2],";	       
	    } elsif ( $bp =~ /^(t|a|u|g|c)/ ) {
		$rna->{'color_r'} .= "0.90,";
		$rna->{'color_g'} .= "0.90,";
		$rna->{'color_b'} .= "0.90,";
	    } else {
		$rna->{'color_r'} .= "1.00,";
		$rna->{'color_g'} .= "1.00,";
		$rna->{'color_b'} .= "1.00,";
	    }
            $i++;
	}
    }

    # add align_bp to pairingmask
    my ( @label, $pal, $pms, $ll, $lab, $gap, $gap2, $a, $end, @alignpos, $middle, $space, @labpos, $here, $count, $there );
    $length = scalar @pm;
    $a = 0;
    $gap = 0;
    @alignpos = split ( /,/, $pms[0]->{'alignpos'} );
    
    foreach $pm ( @pm ) {
        if ( $pm eq "-" ) {
            $pms[0]->{'align_bp'} .= ".,";
            $gap++;
        } else {

            # push positions of current label

            @labpos = ( );
            for ( $i = 0; $i < $length; $i++ ) {  
                if ( $pm[$i] eq $pm ) {
                    push ( @labpos, $alignpos[$i] );
                }
            }

            # where are we in the array?

            $here = $alignpos[$a];
            $count = 0;
            foreach $lab ( @labpos ) {
                if ( $lab == $here ) {
                    last;
                }
                $count++;
            }

            # now get the symmetrical position
            $there = scalar @labpos - $count;
            
            $pms[0]->{'align_bp'} .= "$labpos[$there-1],";
            
        }
        $a++;
    }
    return $entries;
}