source: cpp/frams/canvas/neurodiagram.cpp @ 1298

// This file is a part of Framsticks SDK.  http://www.framsticks.com/
// Copyright (C) 1999-2021  Maciej Komosinski and Szymon Ulatowski.
// See LICENSE.txt for details.

#include "neurodiagram.h"
#include "nn_layout.h"
#include <frams/neuro/neurolibrary.h>
#include <frams/mech/mechworld.h>
#include <frams/util/multirange.h>
#include "canvasutil.h"
#include <frams/neuro/neuroimpl.h>
#include <frams/model/modelobj.h>
#include <frams/simul/simul.h>
#include "common/nonstd_time.h"
#include <algorithm>

#define FIELDSTRUCT NeuroDiagram
ParamEntry neurodiagram_paramtab[] =
{
        { "NeuroDiagram", 1, 4, "NeuroDiagram", "Can be used as the client object in the Window.", },
        { "new", 0, PARAM_USERHIDDEN, "create new NeuroDiagram", "p oNeuroDiagram", PROCEDURE(p_new), },
        { "showCreature", 0, PARAM_USERHIDDEN | PARAM_NOSTATIC, "show dynamic NN", "p(oCreature)", PROCEDURE(p_showcr), },
        { "showModel", 0, PARAM_USERHIDDEN | PARAM_NOSTATIC, "show static NN", "p(oModel)", PROCEDURE(p_showmod), },
        { "hide", 0, PARAM_USERHIDDEN | PARAM_NOSTATIC, "hide NN", "p()", PROCEDURE(p_hide), },
        { 0, 0, 0, },
};
#undef FIELDSTRUCT

Param neurodiagram_param(neurodiagram_paramtab, 0);

static struct ColorDefs colordefs;

void NeuroDiagram::p_new(ExtValue*args, ExtValue*ret)
{
        NeuroDiagram *d = new NeuroDiagram(&colordefs);
        d->drawbackground = false;
        ret->setObject(ExtObject(&neurodiagram_param, d));
}

void NeuroDiagram::p_showcr(ExtValue*args, ExtValue*ret)
{
        Creature *cr = 0;
        if (args->type == TObj)
        {
                const ExtObject& o = args->getObject();
                cr = (Creature*)o.getTarget();
        }
        showLive(cr);
}

void NeuroDiagram::p_showmod(ExtValue*args, ExtValue*ret)
{
        Model *mod = ModelObj::fromObject(args[0]);
        show(mod);
}

static void addNeuroDescription(SString &t, Neuro *n)
{
        static Param par;
        SString c(n->getClassName());
        NeuroClass* cl = n->getClass();
        t += c;
        t += " (";
        if (cl)
                t += cl->getLongName();
        else
                t += "Unknown";
        t += ")";
}

NeuroDiagram::NeuroDiagram(ColorDefs *cd)
        :FramDrawToolkit(cd), livewire(false), indestructor(false), showing_not_alive_label(false), o(0),
        warn_if_not_alive(true), selection(*this), drawbackground(true), linetype(true), layouttype(2)
{
        scroll.setMargin(10, 10); // appropriate size should be adjusted
        dontPaintOutside(0);
        add(&scroll);
        pluginactive = false;
        FramDrawToolkit::setBackColor(ColorDefs::neurobackground);
}

NeuroDiagram::~NeuroDiagram()
{
        indestructor = 1;
        hide();
        remove(&scroll);
        updatePlugin();
}

void NeuroDiagram::hide()
{
        showing_not_alive_label = 0;
        if (o) o->delmodel_list.remove(killnode);
        FOREACH(NeuroProbe*, pr, probes)
                delete pr;
        probes.clear();
        selection.clear(0);
        cr = 0;
        livewire = false;
}

class NNLayoutState_Neurodiagram : public NNLayoutState
{
public:
        NeuroDiagram *nd;
        NNLayoutState_Neurodiagram(NeuroDiagram *_nd) :nd(_nd) {}

        int GetElements()
        {
                return nd->scroll.count();
        }

        int *GetXYWH(int el)
        {
                return &nd->scroll.getInfo(el)->pos.x;
        }

        void SetXYWH(int el, int x, int y, int w, int h)
        {
                ScrollInfo *si = nd->scroll.getInfo(el);
                si->pos.set(x, y); si->size.set(w, h);
        }

        int GetInputs(int el)
        {
                return nd->getNS(el)->n->getInputCount();
        }

        int GetLink(int el, int i)
        {
                return nd->getNS(el)->n->getInput(i)->refno;
        }

        int *GetLinkXY(int el, int i)
        {
                static int XY[2];
                int *xywh = GetXYWH(el);
                XY[0] = 0;
                XY[1] = ((1 + i) * xywh[3]) / (GetInputs(el) + 1);
                return XY;
        }
};


void NeuroDiagram::show(Model *o_)
{
        hide();
        o = o_;
        scroll.removeAll();
        if (o)
        {
                Neuro *n;
                int i;
                killnode = o->delmodel_list.add(STATRICKCALLBACK(this, &NeuroDiagram::onKill, 0));

                // create symbol objects
                for (i = 0; n = o->getNeuro(i); i++)
                {
                        NeuroSymbol *ns = new NeuroSymbol(*this, n);
                        scroll.add(ns, 1); // autodel   
                }
                if (i)
                {
                        struct NNLayoutFunction *nnfun = &nn_layout_functions[layouttype];
                        NNLayoutState_Neurodiagram nn(this);
                        nnfun->doLayout(&nn);
                }
                scroll.invalidate();
                scroll.autoZoom();
        }

        updatePlugin();
        requestPaint();
}

void NeuroDiagram::showLive(Creature *_cr)
{
        showing_not_alive_label = 0;
        if (!_cr) { show(0); return; }
        show(&_cr->getModel());
        cr = _cr;
        livewire = true;
        updatePlugin();
        requestPaint();
}

void NeuroDiagram::paint()
{
        if (drawbackground)
        {
                setColor(ColorDefs::neurobackground);
                clear();
        }

        if (countNeurons() > 0)
        {
                CanvasWindowContainer::paint();
        }
        else
        {
                setColor(ColorDefs::neuroneuron);
                drawAlignedText(size.x / 2, (size.y - textHeight()) / 2, 0, "[No neural network]");
        }

        if (showing_not_alive_label)
        {
                if (time(0) > showing_not_alive_label)
                        showing_not_alive_label = 0;
                else
                {
                        setColor(0, 0, 0);
                        drawAlignedText(not_alive_location.x, not_alive_location.y, 0, "select a creature");
                        drawAlignedText(not_alive_location.x, not_alive_location.y + textHeight(), 0, "to probe neurons");
                }
        }
}

void NeuroDiagram::resize(int w, int h)
{
        CanvasWindowContainer::resize(w, h);
        scroll.autoZoom();
}

int NeuroDiagram::countNeurons()
{
        return scroll.count();
}

void NeuroDiagram::addProbe(int i)
{
        if (i >= countNeurons()) return;
        NeuroProbe *probe = new NeuroProbe(getNS(i));
        Pixel s = getSize();
        s.x = s.y = std::max(probe->getSize().x, std::min(s.x / 3, s.y / 3));
        probes += (void*)probe;
        add(probe);
        probe->resizeClient(s);
        updatePlugin();
        requestPaint(); //because client resized
}

void NeuroDiagram::onKill(void*obj, intptr_t dummy)
{
        show(0);
}

///////////////////////////

NeuroSymbol::NeuroSymbol(NeuroDiagram &nd, Neuro * _n)
        :FramDrawToolkit(nd.getColorDefs()), selected(0), n(_n), diagram(nd)
{
        tooltip = "#";
        tooltip += SString::valueOf((int)n->refno);
        tooltip += " - ";
        label = tooltip;
        addNeuroDescription(tooltip, n);
        label += n->getClassName();
        if (n->getClassParams().length())
        {
                tooltip += "\n"; tooltip += n->getClassParams();
        }
}

void NeuroSymbol::paint()
{
        if (selected)
        {
                setColor(ColorDefs::neuroselection);
                fillRect(0, 0, size.x, size.y);
        }
        diagram.setClip();
        diagram.setColor(ColorDefs::neuroneuron);
        drawNeuroSymbol(this, n->getClass(), 0, 0, size.x, size.y);

        if (size.y > 4 * textHeight())
        {
                const char* t = label.c_str();
                setColor(ColorDefs::neurosymbol);
                drawAlignedText(size.x / 2, size.y - textHeight(), 0, t);

                NeuroImpl *ni = NeuroNetImpl::getImpl(n);
                if (ni && (ni->getChannelCount() > 1))
                {
                        drawLine(size.x - size.x / 16, size.y / 2 - size.y / 16,
                                size.x - size.x / 8, size.y / 2 + size.y / 16);
                        char t[20];
                        sprintf(t, "%d", ni->getChannelCount());
                        moveTo(size.x, size.y / 2 - textHeight());
                        drawText(t);
                }
        }

        /*

                                NeuroDiagram
                                *........................................
                                .                                       .
                                .                    NeuroSymbol        .
                                .      (pos.x,pos.y)-*.........         .
                                .                    . |\     . ^  s    .
                                .            ..._____._| \    . |  i    .
                                .                    . |  \___. |  z    .
                                .                  __._|  /   . |  e    .
                                .                 |  . | /    . |  .    .
                                .                 |  . |/     . |  y    .
                                .                 |  .......... v       .
                                .                 |  <-------->         .
                                .      .......... |    size.x           .
                                .      . |\     . |                     .
                                .    __._| \    . |                     .
                                .      . |  \___._|                     .
                                .    __._|  /   .                       .
                                .   |  . | /    .                       .
                                .   |  . |/     .                       .
                                .   |  ..........                       .
                                .   |                                   .
                                .   |________________...                .
                                .........................................

                                */

                                // NeuroSymbol is also responsible for drawing connection lines from its inputs to other NeuroSymbols' outputs
        NeuroSymbol *ns2;
        if (!diagram.isLive())
                diagram.setColor(ColorDefs::neurolink);
        for (int iw = 0; iw < n->getInputCount(); iw++)
        {
                ns2 = diagram.getNS(n->getInput(iw)->refno); // the other NeuroSymbol (our input will connect to its output)

                int yw = inputY(iw);
                int xw = yw / 4; // x coordinate of the first corner point, depends on yw to avoid overlapping between inputs
                drawLine(size.x / 4, yw, xw, yw); // first horizontal segment (to the left)
                if (diagram.isLive())
                        diagram.setWireColor(ns2->n->state, 0);
                // linetype: 1 (default) - draw straight or U-shape depending on layout
                //           0 (debug option) - only draw straight lines
                if ((diagram.linetype != 1) || (ns2->pos.x + ns2->size.x / 2 < pos.x))
                { // straight line to the other neuron's output (the signal goes forwards)
                        ns2->lineTo(ns2->size.x, ns2->size.y / 2);
                }
                else
                { // make an U-shaped loop from 3 segments - vert/horiz/vert (the signal goes backwards)
                        int y2;
                        int down;
                        if (ns2 == this) down = (iw >= ((n->getInputCount()) / 2)); else down = (ns2->pos.y > (pos.y + size.y));
                        if (down)
                        {
                                y2 = (pos.y + size.y + (size.y - yw) / 3);
                        }
                        else
                        {
                                y2 = pos.y - yw / 3;
                                if ((ns2->pos.y < pos.y) && (ns2->pos.y > (pos.y - ns2->size.y))) y2 -= pos.y - ns2->pos.y;
                        }
                        // note: "diagram" uses global coordinate system, so we add "pos" or "ns2->pos" to get NeuroSymbol's global positions
                        diagram.lineTo(pos.x + xw, y2);
                        diagram.lineTo(ns2->pos.x + ns2->size.x, y2);
                        diagram.lineTo(ns2->pos.x + ns2->size.x, ns2->pos.y + ns2->size.y / 2);
                }

        }
}

void NeuroSymbol::mouse(int x, int y, int t)
{
        if ((t & (LeftButton | ShiftButton)) == (LeftButton | ShiftButton))
        {
                ScrollManager& sc = diagram.scroll;
                sc.setPos2(n->refno, pos.x + x - diagram.symboldragpos.x, pos.y + y - diagram.symboldragpos.y);
                sc.validate();
                requestPaint();
        }
}

int NeuroSymbol::mouseclick(int x, int y, int t)
{
        if ((t & (LeftButton | DblClick)) == (LeftButton | DblClick))
        {
                if (diagram.isLive())
                        diagram.addProbe(n->refno);
                else
                {
                        if (diagram.warn_if_not_alive)
                        {
                                diagram.showing_not_alive_label = time(0) + 10;
                                diagram.not_alive_location.x = pos.x + x;
                                diagram.not_alive_location.y = pos.y + y;
                                diagram.requestPaint();
                        }
                }
                return LeftButton | DblClick;
        }

        if ((t & (LeftButton | ShiftButton)) == (LeftButton | ShiftButton))
        {
                if (selected)
                        diagram.selection.remove(Model::neuroToMap(n->refno));
                else
                        diagram.selection.add(Model::neuroToMap(n->refno));
                diagram.symboldragpos.set(x, y);
                return LeftButton | ShiftButton;
        }

        if (t & LeftButton)
        {
                diagram.selection.set(Model::neuroToMap(n->refno));
                return LeftButton;
        }

        return 0;
}

// coordinate y of i-th input
int NeuroSymbol::inputY(int i)
{
        return (1 + i) * size.y / ((n->getInputCount()) + 1);
}

SString NeuroSymbol::hint(int x, int y)
{
        if ((y >= 0) && (y < size.y))
                if (x < size.x / 4)
                { // inputs?
                        if (n->getInputCount() > 0)
                        {
                                int i = (y * n->getInputCount()) / size.y;
                                double w;
                                Neuro* target = n->getInput(i, w);
                                if (target)
                                {
                                        SString t = "connected to #";
                                        t += SString::valueOf((int)target->refno);
                                        t += " - ";
                                        addNeuroDescription(t, target);
                                        //              if (w!=1.0)
                                        {
                                                t += ", weight=";
                                                t += SString::valueOf(w);
                                        }
                                        return t;
                                }
                        }
                }
        return CanvasWindow::hint(x, y);
}

/////////////////////////

NeuroProbe::NeuroProbe(NeuroSymbol* ns)
        :DCanvasWindow(DCanvasWindow::Title + DCanvasWindow::Border + DCanvasWindow::Close + DCanvasWindow::Size,
                ns->getLabel().c_str(), &neurochart, &neurochart)
{
        holdismine = 0;
        drawing = 0; whichdrawing = -1;
        clientbordersset = 0;
        adjustingvalue = 0;
        link = ns;
        tooltip = SString("Probe for ") + ns->tooltip;
        setPos(ns->getPos().x, ns->getPos().y);
        neurochart.printMinMax(0);
        neurochart.data.setMinMax(-1, 1);
        chnum = 1; chnum2 = 0; chsel = 0;
        chselwidth = 0;
        chselecting = 0;
        updateChannelCount(NeuroNetImpl::getImpl(link->n));
}

void NeuroProbe::onClose()
{
        link->diagram.probes -= this;
        delete this;
}

NeuroProbe::~NeuroProbe()
{
        if (holdismine)
                link->n->flags &= ~Neuro::HoldState;
}

void NeuroProbe::paint()
{
        static char t[40];
        if (!clientbordersset)
        {
                clientbordersset = 1;
                setClientBorder(0, 1, 16, textHeight() + 2); // don't use textheight outside paint/mouse events
        }
        int hold = link->n->flags & Neuro::HoldState;
        float state = (float)link->n->state;
        NeuroImpl *ni = 0;
        if (chsel != 0)
        {
                ni = NeuroNetImpl::getImpl(link->n);
                if (chsel < 0)
                {
                        int dr = -chsel - 1;
                        if (whichdrawing != dr)
                        {
                                drawing = ni->getDrawing(dr);
                                whichdrawing = dr;
                        }
                        if (drawing)
                        {
                                int *dr = drawing;
                                int w = size.x - 2, h = size.y - clienttop - clientbottom;
                                int scale = std::min(w, h);
                                int x0 = clienttop + leftborder + ((w > h) ? (w - h) / 2 : 0);
                                int y0 = clientleft + topborder + ((h > w) ? (h - w) / 2 : 0);

                                while (*dr != NeuroImpl::ENDDRAWING)
                                {
                                        int first = 1;
                                        unsigned char r, g, b;
                                        FramDrawToolkit::splitRGB(*(dr++), r, g, b);
                                        setColor(r, g, b);
                                        while (*dr != NeuroImpl::ENDDRAWING)
                                        {
                                                int x = ((*(dr++)) * scale) / (NeuroImpl::MAXDRAWINGXY + 1) + x0;
                                                int y = ((*(dr++)) * scale) / (NeuroImpl::MAXDRAWINGXY + 1) + y0;
                                                if (first) { moveTo(x, y); first = 0; }
                                                else lineTo(x, y);
                                        }
                                        dr++;
                                }
                        }
                }
        }
        DCanvasWindow::paintWithClient((chsel < 0) ? 0 : client);
        setColor(0, 0, 0);
        int yline = size.y - 2;
        if (chsel >= 0)
        {
                if (ni) state = (float)ni->getState(chsel);
                yline -= textHeight();
                int y = mapClientY(neurochart.mapData(state));
                int x = size.x - 15 - 1;
                drawLine(1, yline, size.x - 2, yline);
                if (hold)
                {
                        sprintf(t, "hold: %1.3g", state);
                        fillRect(x, y - 1 - 5, 15, 3 + 5 + 5);
                        setColor(255, 0, 0);
                        fillRect(x + 2, y - 1, 15 - 2 - 2, 3);
                }
                else
                {
                        sprintf(t, "signal: %1.3g", state);
                        fillRect(x, y - 1, 15, 3);
                }
                drawAlignedText(size.x - textHeight(), yline, 1, t);
        }

        if ((chnum > 1) || (chnum2 > 0))
        {
                if (chselecting) setColor(255, 255, 255); else setColor(0, 70, 0);
                if (chsel < 0)
                        sprintf(t, "%c/%c", 'A' - chsel - 1, 'A' + chnum2 - 1);
                else
                        sprintf(t, "%d/%d", chsel, chnum);
                moveTo(0, yline - textHeight());
                chselwidth = textWidth(t);
                drawText(t, -1, getSize().x);
        }
        else
                chselwidth = 0;
}

void NeuroProbe::mouse(int x, int y, int b)
{
        if (chselecting)
        {
                int ch = chsel0 + (x - chselx0) / 10;
                if (selectChannel(ch)) requestPaint();
                b &= ~LeftButton;
        }
        DCanvasWindow::mouse(x, y, b);
        if (adjustingvalue)
        {
                double st = neurochart.unmapData(unmapClientY(y));
                if (st < -1.0) st = -1.0; else if (st > 1.0) st = 1.0;
                if (chsel == 0)
                        link->n->state = st;
                else if (chsel >= 0)
                {
                        NeuroImpl *ni = NeuroNetImpl::getImpl(link->n);
                        if (ni) ni->setCurrentState(st, chsel);
                }
                requestPaint();
        }
}

void NeuroProbe::mouseunclick(int x, int y, int b)
{
        adjustingvalue = 0;
        chselecting = 0;
        DCanvasWindow::mouseunclick(x, y, b);
}

bool NeuroProbe::insideChSelector(int x, int y)
{
        if ((x > 0) && (x < chselwidth))
        {
                int sy = size.y;
                if (chsel >= 0) sy -= textHeight();
                return ((y < sy) && (y > (sy - textHeight())));
        }
        return 0;
}

int NeuroProbe::mouseclick(int x, int y, int b)
{
        if ((b & LeftButton) && insideChSelector(x, y))
        {
                chselx0 = x; chsel0 = chsel;
                chselecting = 1;
                requestPaint();
                return LeftButton;
        }
        int ret = DCanvasWindow::mouseclick(x, y, b);
        if (ret)
        {
                link->diagram.selection.set(Model::neuroToMap(link->n->refno));
                return ret;
        }
        if (b & LeftButton)
        {
                if (x > size.x - 16)
                {
                        link->n->flags |= Neuro::HoldState;
                        holdismine = 1;
                        adjustingvalue = 1;
                        mouse(x, y, b);
                        return LeftButton;
                }
                else if (y > size.y - 16)
                {
                        link->n->flags ^= Neuro::HoldState;
                        holdismine = ((link->n->flags & Neuro::HoldState) != 0);
                        requestPaint();
                        return LeftButton;
                }
        }
        return 0;
}

SString NeuroProbe::hint(int x, int y)
{
        if ((chsel >= 0) && (x < size.x - 16) && (y > size.y - 16))
                return SString((link->n->flags & Neuro::HoldState) ? "Click to release" : "Click to hold");
        else if (insideChSelector(x, y))
                return SString::sprintf("channel %d of %d (click and drag to switch channels)", chsel, chnum);
        return DCanvasWindow::hint(x, y);
}

/** @return true == channel changed */
bool NeuroProbe::selectChannel(int ch)
{
        if (ch < -chnum2) ch = -chnum2; else if (ch >= chnum) ch = chnum - 1;
        if (ch == chsel) return false;
        chsel = ch;
        neurochart.data.clear();
        return true;
}

void NeuroProbe::updateChannelCount(NeuroImpl *ni)
{
        if (!ni) return;
        chnum = ni->getChannelCount();
        chnum2 = ni->getDrawingCount();
        if (chsel >= chnum) selectChannel(chnum - 1);
        if (chsel < -chnum2) selectChannel(-chnum2);
}

void NeuroProbe::sampling()
{
        NeuroImpl *ni = NeuroNetImpl::getImpl(link->n);
        updateChannelCount(ni);
        if (!chsel)
                neurochart.data += (float)(link->n->state);
        else
                neurochart.data += (float)(ni->getState(chsel));
        whichdrawing = -1;
}

////

void NeuroDiagram::probeSampling(void*obj, intptr_t dummy)
{
        FOREACH(NeuroProbe*, pr, probes) pr->sampling();
        requestPaint();
}

void NeuroDiagram::updatePlugin()
{
        //int needplugin=(!probes)>0;
        bool needplugin = livewire;
        if (needplugin == pluginactive) return;
        if (needplugin)
        {
                if (!cr) return;
                sim = cr->group->getLibrary().sim;
                pluginnode = sim->l_plugin.add(STATRICKCALLBACK(this, &NeuroDiagram::probeSampling, 0));
        }
        else
                sim->l_plugin.remove(pluginnode);
        pluginactive = needplugin;
}

/////////////

void NeuroDiagramSelection::updateSelection(MultiRange& newsel)
{
        MultiRange added = getAdded(newsel);
        if (!added.isEmpty())
        {
                added.shift(Model::mapToNeuro(0));
                added.intersect(0, diagram.countNeurons() - 1);
                for (int i = 0; i < added.rangeCount(); i++)
                {
                        const IRange &r = added.getRange(i);
                        for (int j = r.begin; j <= r.end; j++)
                                diagram.getNS(j)->selected = 1;
                }
        }
        MultiRange removed = getRemoved(newsel);
        if (!removed.isEmpty())
        {
                removed.shift(Model::mapToNeuro(0));
                removed.intersect(0, diagram.countNeurons() - 1);
                for (int i = 0; i < removed.rangeCount(); i++)
                {
                        const IRange &r = removed.getRange(i);
                        for (int j = r.begin; j <= r.end; j++)
                                diagram.getNS(j)->selected = 0;
                }
        }
        if (!diagram.indestructor) diagram.requestPaint();
}