from rdkit import Chem
from rdkit.Chem import Draw, rdMolTransforms
import ipywidgets as widgets
from traitlets import Unicode, Int, validate
class MolSVGWidget(widgets.DOMWidget):
    _view_name = Unicode('MolSVGView').tag(sync=True)
    _view_module = Unicode('molsvg_widget').tag(sync=True)
    _view_module_version = Unicode('0.0.1').tag(sync=True)
    
    svg = Unicode('', help="svg to be rendered").tag(sync=True)
    selected_atoms = Unicode('', help="list of currently selected atoms").tag(sync=True)

%%javascript
// make sure our module is only defined
// only once.
require.undef('molsvg_widget');

// Define the `molsvg_widget` module using the Jupyter widgets framework.
define('molsvg_widget', ["@jupyter-widgets/base"],
       function(widgets) {

    // The frontend class:
    var MolSVGView = widgets.DOMWidgetView.extend({

        // This method creates the HTML widget.
        render: function() {
            this.svg_div = document.createElement('div');
            this.el.appendChild(this.svg_div);
            this.model.on('change:svg', this.svg_changed, this);
            this.svg_changed();
        },
        
        // called when the SVG is updated on the Python side
        svg_changed: function() {
            var txt = this.model.get('svg'); 
            this.svg_div.innerHTML = txt;
            var sels = this.svg_div.getElementsByClassName("atom-selector");
            for(var i=0;i<sels.length;i++){
                sels[i].onclick = (evt) => { return this.atom_clicked(evt) };
            }
            
        },

        // callback for when an atom is clicked
        atom_clicked: function(evt) {
            //alert("  "+evt+"|"+this);
            if(!evt.currentTarget.getAttribute('class')){
                return;
            }
            var satmid = evt.currentTarget.getAttribute('class').match(/atom-([0-9]+)/);
            if(satmid.length >1){
                var atmid = Number(satmid[1]);
                var curSel = this.model.get('selected_atoms');
                var splitSel = curSel.split(',');
                var selItms = [];
                var idx = -1;
                //alert("|"+atmid+"|"+curSel+"|len: "+splitSel.length);
                if(curSel != "" && splitSel.length>0){
                    selItms = Array.from(splitSel).map(item => Number(item));
                    idx = selItms.indexOf(atmid);
                }
                if(idx == -1){
                    selItms = selItms.concat(atmid);
                    evt.currentTarget.style["stroke-width"]=3;
                    evt.currentTarget.style["stroke-opacity"]=1;
                    evt.currentTarget.style["stroke"]='#AA22FF';
                } else {
                    selItms.splice(idx,1);
                    evt.currentTarget.style["stroke-width"]=1;
                    evt.currentTarget.style["stroke-opacity"]=0;
                    evt.currentTarget.style["stroke"]='#FFFFFF';
                }
                this.model.set('selected_atoms',String(selItms));
                this.touch();
            }
        }

    });

    return {
        MolSVGView : MolSVGView
    };
});

def mol_to_svg(mol):
    d = Draw.MolDraw2DSVG(350, 200)
    d.SetFontSize(0.8)
    d.DrawMolecule(mol)
    d.TagAtoms(mol)
    d.FinishDrawing()
    return d.GetDrawingText()

bilastine_smi = "O=C(O)C(c1ccc(cc1)CCN4CCC(c2nc3ccccc3n2CCOCC)CC4)(C)C"

bilastine = Chem.MolFromSmiles(bilastine_smi)
dbilastine = Draw.PrepareMolForDrawing(bilastine)

w = MolSVGWidget(svg=mol_to_svg(dbilastine))

w

w.selected_atoms

'10,11,12,13'

selected_i = [int(i) for i in w.selected_atoms.split(',')]

selected_i

[10, 11, 12, 13]

rdMolTransforms.SetDihedralDeg(dbilastine.GetConformer(), *selected_i,
    rdMolTransforms.GetDihedralDeg(dbilastine.GetConformer(), *selected_i) + 180)

w = MolSVGWidget(svg=mol_to_svg(dbilastine))

w