#!/usr/bin/perl

use strict;
use warnings;

#
# jmol2pymol.sh v. 1.0 [2018-12-20]
# d.a.suplatov@belozersky.msu.ru
#

#THIS SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND
#PLEASE CHECK MANUALLY THE CONSISTENCY OF THE GENERATED OUTPUT
#PLEASE REPORT ANY PROBLEMS TO D.A.SUPLATOV@BELOZERSKY.MSU.RU

if (@ARGV != 3) {
    print "\n";
    print "This script will convert the JMOL instruction file xxx.spt into a PyMOL instruction file\n";
    print "The JMOL instruction file xxx.spt will be automatically produced by parMATT/MATT software\n";
    print "Usage: $0 <parMATT_output.spt> <parMATT_output.txt> <output.py>\n";
    print "          <parMATT_output.spt> - the JMOL instruction file xxx.spt automatically produced by parMATT/MATT\n";
    print "          <parMATT_output.txt> - the plain-text information file automatically produced by parMATT/MATT\n";
    print "                <output.pymol> - the name of the output file with instructions for PyMOL\n";
    print "\n";
    exit 1;
}

print "\n";
print "THIS SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND\n";
print "PLEASE CHECK MANUALLY THE CONSISTENCY OF THE GENERATED OUTPUT\n";
print "PLEASE REPORT ANY PROBLEMS TO D.A.SUPLATOV\@BELOZERSKY.MSU.RU\n";
print "\n";

my $INPUT = $ARGV[0];
my $INFO = $ARGV[1];
my $OUTPUT = $ARGV[2];

if ( ! -e $INPUT ) {
    print "Error: The input file $INPUT does not exist\n";
    exit 1;
}

if ( -e $OUTPUT ) {
    print "Error: The output file $OUTPUT already exists\n";
    exit 1;
}

my %data = ();
print "Info: Reading the information file $INFO\n";
open (I, $INFO) || die "Error: Failed to open file $INFO for reading\n";
while (<I>) {
    chomp;
    
    if (/^(.*):(.*)\s+[-A-Za-z]+\s+\d+\s+\((.*)\)/) { 
        my $pdb = $1; 
        my $pdb_chain = $2;        
        my $out_chain = $3;
        
        $pdb =~ s/\s+//g;
        $pdb_chain =~ s/\s+//g;
        $out_chain =~ s/\s+//g;
        
        my $name = $pdb."_$pdb_chain";
        
        if ($data{$name} && $out_chain ne $data{$name}) {
            print "Error: PDB $name is assigned to multiple chains ($data{$name}, $out_chain) in the $INFO file\n";
            print "Error: This script will terminate here\n";
            print "Error: Check manually the consistency of your data and contact the Support if necessery\n";
            exit 1;
        }
        else {$data{$name} = $out_chain;}
    }
}
close(I);

my @pdbs = keys %data;
print "Info: Information on ".(scalar @pdbs)." PDB IDs has been collected\n";

my %instr = ();
print "Info: Reading the JMOL annotation file $INPUT\n";
open (FI, $INPUT) || die "Error: Failed to open file $INPUT for reading\n";
while (<FI>) {
    if (/^select all/) {next;}
    if (/^select/) {
        chomp;
        my $line = $_;
        $line =~ s/select\s+//;
        my @arr = split(/,/, $line);
        foreach my $elm (@arr) {
            if ($elm =~ /(\d+-\d+):(.*)/) {
                my $resid = $1; 
                my $out_chain = $2;
                $resid =~ s/\s+//g;
                $out_chain =~ s/\s+//g;
                push(@{$instr{$out_chain}}, $resid);
            }
            else {
                print "Warning: The residue specification \"$elm\" in data line \"$line\" will be skipped\n";
            }
        }
    }
}
print "Info: Instructions have been collected\n";
close(FI);;

print "Info: Reformatting the JMOL instructions into the PyMOL format and printing the output file\n";
open (P, ">$OUTPUT") || die "Error: Failed to open file $OUTPUT for writing\n";
print P "#INSTRUCTIONS\n";
print P "#\n";
print P "#First, separate the parMATT_output.pdb file with superimposed coordinates into individual PDB files\n";
print P "#For this you should use the splitMATT2chains.sh script (see https://biokinet.belozersky.msu.ru/parmatt, \"Download accessory scripts\")\n";
print P "#\n";
print P "#Load all superimposed PDB files into PyMol, and then run this script via the PyMol console:\n";
print P "#> run path/to/this/script/instructions.py\n";
print P "#\n";
print P "#Finally, you may wish to create a high-quality illustration from what you see in the 3D-viewer\n";
print P "#Run these commands in the PyMol console (NOTE: you may need to adjust the X Y resolution):\n";
print P "#set ambient, 0.5\n";
print P "#ray 1280, 1024\n";
print P "#\n";
print P "cmd.hide(\"everything\")\n";
print P "cmd.bg_color(\"white\")\n";
print P "cmd.set(\"cartoon_discrete_colors\", \"1\")\n";
print P "cmd.set(\"cartoon_oval_length\", \"0.3\")\n";
print P "cmd.set(\"cartoon_oval_width\", \"0.035\")\n";
print P "cmd.set(\"cartoon_loop_radius\", \"0.1\")\n";
print P "cmd.set(\"cartoon_rect_width\", \"0.035\")\n";
print P "cmd.set(\"cartoon_rect_length\", \"0.5\")\n";
print P "cmd.show(\"cartoon\")\n";
print P "cmd.color(\"black\")\n";
print P "cmd.color(\"violet\", \"ss h\")\n";
print P "cmd.color(\"yellow\", \"ss s\")\n";
print P "cmd.set(\"ray_shadows\", \"0\")\n";

foreach my $pdb (@pdbs) {
    my $out_chain = $data{$pdb};
    if (!$instr{$out_chain}) {
        print "Error: Failed to collect JMOL instructions for the PDB $pdb which was mentioned in the information file as chain $out_chain\n";
        print "Error: This script will terminate here\n";
        print "Error: Check manually the consistency of your data and contact the Support if necessery\n";
        exit 1;
    }
    my $out = "cmd.color(\"white\", \"$pdb and not resid ";
    foreach my $i (0..$#{$instr{$out_chain}}) {
        my $resid = $instr{$out_chain}->[$i];
        $out .= "$resid";
        if ($i != $#{$instr{$out_chain}}) {$out .= "+";}
    }
    $out .= "\")\n";
    print P $out;
}

print P "#EOF\n";
close (P);

print "Done!\n";
print "\n";
exit 0;

