#!/usr/bin/perl

use strict;
use warnings;

use version;
use Algorithm::Combinatorics qw( variations );
use Chemistry::Atom qw( angle_deg dihedral_deg );
use Chemistry::File::SDF;
use Chemistry::Mol;
use Chemistry::OpenSMILES qw(
    %normal_valence
    clean_chiral_centers
    is_chiral
    is_chiral_planar
    is_triple_bond
    mirror
);
use Chemistry::OpenSMILES qw( is_double_bond );
use Chemistry::OpenSMILES::Stereo qw( chirality_to_pseudograph mark_all_double_bonds );
use Chemistry::OpenSMILES::Writer qw( write_SMILES );
use File::Basename qw( basename );
use Getopt::Long::Descriptive;
use Graph::Nauty qw( orbits );
use Graph::Traversal::DFS;
use Graph::Undirected;
use List::Util qw( all any first max min none sum sum0 );
use SmilesScripts::Chirality qw( chiral_volume square_planar trigonal_bipyramidal );

my $default_planarity_threshold = 1e-6;

my %bond_order = (
    '-' => 1,
    '=' => 2,
    '#' => 3,
    '$' => 4,
);

# Extracted from symmetry.hpp file of the gemmi project. For more information see: 
#
# https://github.com/project-gemmi/gemmi/blob/2d25148dc3ed39adc19422f64be76d7ba6dca2f4/include/gemmi/symmetry.hpp#L912

my @sohncke_groups = qw(
      1   3   4   5  16  17  18  19  20  21  22  23  24  75  76  77  78
     79  80  89  90  91  92  93  94  95  96  97  98 143 144 145 146 149
    150 151 152 153 154 155 168 169 170 171 172 173 177 178 179 180 181
    182 195 196 197 198 199 207 208 209 210 211 212 213 214
);

my $basename = basename $0;
my( $opt, $usage ) = describe_options( <<"END" . 'OPTIONS',
USAGE
    $basename [<args>] [<files>]

DESCRIPTION
    $basename converts SDF files into SMILES.
END
    [ 'convert-type-8-bonds-to-covalent',
      'convert all type 8 bonds to single covalent bonds' ],
    [ 'no-adjust-charges',
      'do not adjust atom charges when converting type 8 bonds' ],
    [ 'no-estimate-attached-hydrogens',
      'do not estimate the number of attached hydrogens' ],
    [],
    [ tetrahedral_chiral_method => hidden => {
        one_of => [
            [ 'ignore-flat-tetrahedra' =>
                'when considering tetrahedral chiral centers, ignore ' .
                'flat tetrahedra (based on volume threshold) [default]' ],
            [ 'require-tetrahedral-angles' =>
                'when considering tetrahedral chiral centers, ignore ' .
                'tetrahedra with average bond angle far from 109.5 degrees' ],
        ],
        default => 'ignore_flat_tetrahedra'
      }
    ],
    [ 'planarity-threshold=f',
      'planarity volume threshold in cubic angstroms; tetrahedra with ' .
      'chiral volume less or equal to the given one will be understood ' .
      "as flat [default: $default_planarity_threshold]",
      { default => $default_planarity_threshold } ],
    [ 'normalise-vectors',
      'normalise vectors used for chiral volume calculation' ],
    [],
    [ 'mark-square-planar-centers',
      'mark chirality in square planar centers' ],
    [],
    [ 'no-generate-racemates',
      'do not attempt to fully represent racemates in non-Sohncke space groups' ],
    [],
    [ 'remove-excess-stereochemistry',
      'remove excess stereochemistry settings (experimental)' ],
    [],
    [ 'help', 'print usage message and exit', { shortcircuit => 1 } ],
);

if( $opt->help ) {
    print $usage->text;
    exit;
}

@ARGV = ( '-' ) unless @ARGV;

Chemistry::Mol->register_format( sdf => Chemistry::File::SDF:: );

# Perl Clone module does native cloning, thus it should be able to handle
# larger molecules than the default Storable cloner.
$Chemistry::Mol::clone_backend = 'Clone';

if( $opt->remove_excess_stereochemistry &&
    version->parse( $Chemistry::OpenSMILES::VERSION ) < version->parse( '0.8.6' ) ) {
    print STDERR "${basename}:: WARNING, Chemistry::OpenSMILES versions before 0.8.6 " .
                 'remove chiral markers from anomers.', "\n";
}

foreach my $filename (@ARGV) {
    local $SIG{__WARN__} = sub {
        print STDERR "$basename: $filename: $_[0]"
    };

    eval {
        my $reader = Chemistry::Mol->file( $filename, format => 'sdf' );
        $reader->open( '<' );
        while( my $molecule = $reader->read_mol( $reader->fh ) ) {
            # Convert [C-]#[O+] -> [C]#[O]
            # See https://projects.ibt.lt/repositories/issues/1621 for justification
            my @atoms_without_H;
            for my $bond ($molecule->bonds) {
                next unless $bond->type == 3;
                my( $C, $O ) = sort { $a->symbol cmp $b->symbol } $bond->atoms;
                next unless $C->symbol eq 'C';
                next unless $O->symbol eq 'O';
                next unless $C->formal_charge == -1;
                next unless $O->formal_charge ==  1;
                next unless all { $_->type == 8 } ( $C->bonds($O), $O->bonds($C) );
                $C->formal_charge(0);
                $O->formal_charge(0);
                push @atoms_without_H, $C;
            }

            # Resolving the issue with bonds of order 8
            my @type_8_bonds;
            for my $bond ($molecule->bonds) {
                next unless $bond->type == 8;
                push @type_8_bonds, $bond;

                if( $opt->convert_type_8_bonds_to_covalent ) {
                    $bond->type(1);
                }

                if( grep { !$_->formal_charge } $bond->atoms ) {
                    $bond->type(1) if !$opt->convert_type_8_bonds_to_covalent;
                    next;
                }
                
                my @atoms = sort { $a->formal_charge <=>
                                   $b->formal_charge } $bond->atoms;
                next unless $atoms[0]->formal_charge < 0 &&
                            $atoms[1]->formal_charge > 0;
                my $min_abs_charge = min -$atoms[0]->formal_charge,
                                          $atoms[1]->formal_charge;

                if( $opt->convert_type_8_bonds_to_covalent ) {
                    next if $opt->no_adjust_charges;
                    # When bond 8 is converted into plain covalent bond,
                    # charges for both atoms have to be adjusted
                    $atoms[0]->formal_charge( $atoms[0]->formal_charge +
                                              $min_abs_charge );
                    $atoms[1]->formal_charge( $atoms[1]->formal_charge -
                                              $min_abs_charge );
                } else {
                    # Replace bond 8 with a single bond if both of the atoms
                    # have opposite charges and decrease the charges by 1.
                    # Further increase the bond order if both of the atoms
                    # still have opposing charges.
                    # TODO: Not sure what to do with bond orders higher than 3.
                    my $bond_type = $bond->type == 8 ? 0 : $bond->type;
                    my $bond_type_new = min $bond_type + $min_abs_charge, 3;
                    next if $bond_type >= $bond_type_new;

                    $bond->type( $bond_type_new );
                    next if $opt->no_adjust_charges;
                    $atoms[0]->formal_charge( $atoms[0]->formal_charge +
                                              ($bond_type_new - $bond_type) );
                    $atoms[1]->formal_charge( $atoms[1]->formal_charge -
                                              ($bond_type_new - $bond_type) );
                    # TODO: What to do when both charges are of the same sign?
                }
            }

            my $name = $molecule->name;
            my @graphs;
            for my $moiety ($molecule->separate) {
                my %atoms;
                my %atom_by_vertex;
                my $graph = Graph::Undirected->new( refvertexed => 1 );
                my @atoms = $moiety->atoms;

                # First pass through atoms to add them to molecular graph 
                for my $i (0..$#atoms) {
                    my $atom = $atoms[$i];
                    my $vertex = {
                        charge => $atom->formal_charge,
                        number => $i,
                        symbol => $atom->symbol,
                    };
                    my $most_common_isotope =
                        Chemistry::Atom->new( symbol => $atom->symbol );
                    if( $most_common_isotope->mass != $atom->mass ) {
                        $vertex->{isotope} = sprintf( '%.0f', $atom->mass ) + 0;
                    }
                    if( any { $_ eq $atom } @atoms_without_H ) {
                        $vertex->{hcount} = 0;
                    }
                    $graph->add_vertex( $vertex );
                    $atoms{$atom} = $vertex;
                    $atom_by_vertex{$vertex} = $atom;
                }

                # Second pass through atoms to set chirality
                for my $i (0..$#atoms) {
                    my $atom   = $atoms[$i];
                    my $vertex = $atoms{$atom};
                    my @neighbours = $atom->neighbors;
                    my $hydrogens  = grep { $_->symbol eq 'H' &&
                                            scalar $_->neighbors == 1 }
                                          @neighbours;

                    if(      @neighbours == 4 && $hydrogens <= 1 ) {
                        # Tetrahedral chiral and square planar centers
                        my $volume = chiral_volume( @neighbours, $opt->normalise_vectors );
                        if( $opt->tetrahedral_chiral_method eq 'ignore_flat_tetrahedra' ) {
                            if( abs $volume > $opt->planarity_threshold ) {
                                $vertex->{chirality} = $volume < 0 ? '@' : '@@';
                            } else {
                                next unless $opt->mark_square_planar_centers;
                                $vertex->{chirality} =
                                    square_planar( map { [ $_->coords->array ] } $atom, @neighbours );
                            }
                        } else {
                            my $angle = sum angle_deg( $neighbours[0], $atom, $neighbours[1] ),
                                            angle_deg( $neighbours[0], $atom, $neighbours[2] ),
                                            angle_deg( $neighbours[0], $atom, $neighbours[3] ),
                                            angle_deg( $neighbours[1], $atom, $neighbours[2] ),
                                            angle_deg( $neighbours[1], $atom, $neighbours[3] ),
                                            angle_deg( $neighbours[2], $atom, $neighbours[3] );
                            if( abs( $angle / 6 - 109.5 ) <= 5 ) {
                                $vertex->{chirality} = $volume < 0 ? '@' : '@@';
                            } else {
                                next unless $opt->mark_square_planar_centers;
                                $vertex->{chirality} =
                                    square_planar( map { [ $_->coords->array ] } $atom, @neighbours );
                            }
                        }
                        $vertex->{chirality_neighbours} =
                            [ map { $atoms{$_} } @neighbours ];
                    } elsif( @neighbours == 5 && $hydrogens <= 1 ) {
                        # Trigonal bipyramidal centers
                        my @pairs = variations( \@neighbours, 2 );
                        my @angles = map { angle_deg( $_->[0], $atom, $_->[1] ) } @pairs;
                        @pairs =  map  { $pairs[$_] }
                                  sort { $angles[$b] <=> $angles[$a] } 0..$#pairs;
                        @angles = sort { $b <=> $a } @angles;
                        next unless $angles[0] > 170; # largest angle should be close to 180
                        # TODO: Check for planarity maybe
                        next; # TODO: Implement
                        $vertex->{chirality} =
                            trigonal_bipyramidal( map { [ $_->coords->array ] } $atom, @neighbours );
                        $vertex->{chirality_neighbours} =
                            [ map { $atoms{$_} }
                                  $pairs[0]->[0],
                                  ( grep { $_ != $pairs[0]->[0] &&
                                           $_ != $pairs[0]->[1] } @neighbours ),
                                  $pairs[0]->[1] ];
                    }
                }

                # First pass through bonds to set them as given in SDF
                for my $bond ($moiety->bonds) {
                    my $type = bond_type( $bond->type );
                    if( $type ) {
                        $graph->set_edge_attribute(
                            map( { $atoms{$_} } $bond->atoms ),
                            'bond',
                            $type );
                    } else {
                        $graph->add_edge(
                            map { $atoms{$_} } $bond->atoms );
                    }
                }

                # Second pass through bonds to find cis/trans bonds
                my @cis_trans_bonds;
                for my $bond ($moiety->bonds) {
                    my @atoms = map { $atoms{$_} } $bond->atoms;
                    next unless is_double_bond( $graph, @atoms );

                    my @atoms1 = sort { $a->{number} <=> $b->{number} }
                                 grep { $_ != $atoms[1] }
                                      $graph->neighbours( $atoms[0] );
                    my @atoms4 = sort { $a->{number} <=> $b->{number} }
                                 grep { $_ != $atoms[0] }
                                      $graph->neighbours( $atoms[1] );
                    next unless @atoms1 && @atoms4;
                    next if @atoms1 > 2 || @atoms4 > 2;

                    @atoms = ( shift @atoms1, @atoms, shift @atoms4 );
                    my $angle = dihedral_deg( map { $atom_by_vertex{$_} } @atoms );
                    push @cis_trans_bonds, [ @atoms, abs $angle < 90 ? 'cis' : 'trans' ];
                }

                # Third pass through atoms to add the attached hydrogens
                if( exists $molecule->attr('sdf/data')->{COD_SDF_ATTACHED_HYDROGEN_ATOMS} ) {
                    my @atoms_in_molecule = $molecule->atoms;
                    for (@{$molecule->attr('sdf/data')->{COD_SDF_ATTACHED_HYDROGEN_ATOMS}}) {
                        my( $atom_id, $implicit_hydrogens ) = split /\s+/, $_;
                        next unless any { $_ eq $atoms_in_molecule[$atom_id-1] } @atoms;
                        $atoms{$atoms_in_molecule[$atom_id-1]}->{hcount} = int $implicit_hydrogens;
                    }
                }

                # Fourth pass through atoms to preserve unusual valencies.
                # This should be done by Chemistry::OpenSMILES, but is not
                # (https://github.com/merkys/Chemistry-OpenSMILES/issues/6).
                # The current workaround is rather hackish: atoms that have
                # unusual valences are marked as having SMILES class 1,
                # which is later removed with a simple regular expression.
                for my $atom ($graph->vertices) {
                    next unless exists $normal_valence{$atom->{symbol}};
                    my $charge = $atom->{charge} ? $atom->{charge} : 0;
                    my $valence = -$charge +
                                  ($atom_by_vertex{$atom}->formal_radical &&
                                   $atom_by_vertex{$atom}->formal_radical > 1
                                   ? $atom_by_vertex{$atom}->formal_radical - 1 : 0);
                    my $participates_in_8_bond;
                    for my $neighbour ($graph->neighbors( $atom )) {
                        # If charges were not adjusted and this was an order 8
                        # bond, it should not be included in valence counting.
                        if( $opt->no_adjust_charges &&
                            grep { $_->[0] eq $atom_by_vertex{$atom} ||
                                   $_->[1] eq $atom_by_vertex{$atom} }
                            grep { $_->[0] eq $atom_by_vertex{$neighbour} ||
                                   $_->[1] eq $atom_by_vertex{$neighbour} }
                                 map { [ $_->atoms ] } @type_8_bonds ) {
                            $participates_in_8_bond = 1;
                            next;
                        }
                        if( $graph->has_edge_attribute( $atom, $neighbour, 'bond' ) ) {
                            $valence += $bond_order{$graph->get_edge_attribute( $atom, $neighbour, 'bond' )};
                        } else {
                            $valence++;
                        }
                    }
                    my $normal_valence = first { $_ >= $valence }
                                               sort @{$normal_valence{$atom->{symbol}}};
                    # In SDF files generated in 'molecules-in-COD' project,
                    # positively charged C atoms usually denote the additional
                    # positive charge in aromatic rings. Therefore, H atoms
                    # are not added for such atoms.
                    if( !$opt->no_estimate_attached_hydrogens &&
                        $normal_valence &&
                        ( $atom->{symbol} ne 'C' || $charge <= 0 ) &&
                        !exists $atom->{hcount} ) {
                        $atom->{hcount} = $normal_valence - $valence;
                    }
                    if( !$normal_valence ||
                        $participates_in_8_bond ||
                        ($atom_by_vertex{$atom}->formal_radical && $atom_by_vertex{$atom}->formal_radical > 1) ) {
                        $atom->{class} = 1;
                    }
                }

                # Remove superfluous cis/trans and chirality settings (experimental)
                if( $opt->remove_excess_stereochemistry ) {
                    my $copy = $graph->copy;
                    chirality_to_pseudograph( $copy );

                    my @orbits = orbits( $copy,
                                         sub { ref $_[0] && exists $_[0]->{symbol} ? &write_SMILES : '' } );

                    my $color_sub = sub {
                                            for my $i (0..$#orbits) {
                                                return $i if any { $_ == $_[0] } @{$orbits[$i]};
                                            }
                                        };

                    # WARNING: The following code does not work as expected with Chemistry::OpenSMILES v0.8.5 or earlier.
                    # There is a safeguard which throws a warning in this script.
                    mark_all_double_bonds( $graph, \@cis_trans_bonds, undef, $color_sub );

                    clean_chiral_centers( $graph, $color_sub );
                } else {
                    mark_all_double_bonds( $graph, \@cis_trans_bonds );
                }

                # Detect chiral molecular entities in non-Sohncke space groups.
                # If a molecular entity has a single chiral center in a non-Sohncke space groups, it is removed.
                # Otherwise the molecule is duplicated with chiral marks inverted.
                # Chiral planar centers are ignored as they are immune to mirror transformation.
                if( !$opt->no_generate_racemates &&
                    exists   $molecule->attr('sdf/data')->{COD_SDF_SPACE_GROUP_IT_NUMBER} &&
                    ( none { $molecule->attr('sdf/data')->{COD_SDF_SPACE_GROUP_IT_NUMBER} == $_ } @sohncke_groups ) &&
                    ( any  { is_chiral $_ && !is_chiral_planar $_ } $graph->vertices ) ) {
                    my $number_of_chiral_atoms = grep { is_chiral $_ && !is_chiral_planar $_ }
                                                      $graph->vertices;
                    push @graphs, $graph if $number_of_chiral_atoms > 1;

                    my $graph = $graph->deep_copy;
                    for my $atom ($graph->vertices) {
                        next unless is_chiral $atom;
                        next if is_chiral_planar $atom;
                        if( $number_of_chiral_atoms == 1 ) {
                            delete $atom->{chirality};
                            delete $atom->{chirality_neighbours};
                        } else {
                            mirror $atom;
                        }
                    }
                    push @graphs, $graph;
                } else {
                    push @graphs, $graph;
                }
            }

            print join( '.', map { $_ = write_SMILES( $_ ); s/:1\]/\]/g; $_ } @graphs ),
                  "\t",
                  $name,
                  "\n";
        }
    };
    if( $@ ) {
        $@ =~ s/\n$//;
        print STDERR "$basename: $filename: $@\n";
    }
}

sub bond_type
{
    my( $type ) = @_;

    return '=' if $type == 2;
    return '#' if $type == 3;
    return ':' if $type == 4;

    return undef; # TODO: discuss if this is correct
}