with(plottools);
loadlib("TensMaps");
loadlib("ColMaps");

DrawTools := module()
option package;
export drawpng,imrng,drawline,drawdot,solidcirc,drawpoly,subcolor,drawsublev,disprng,animrng,drawplot,drawvec,drawbelow;
local topng;

    topng0 := proc(im1::Array(datatype=float[8]),im2::Array(datatype=float[8]),m::integer[4],n::integer[4])
        for i from 1 to m do
            for j from 1 to n do
                if(im1[i,j,1]>1) then
                    im2[i,j,4] := 1.0;
                else
                    for k from 1 to 3 do
                        im2[i,j,k] := im1[i,j,k];
                    end do;
                    im2[i,j,k] := 0.0;
                end if;
            end do;
        end do;
    end proc;

    topng0 := Compiler:-Compile(topng0);

    topng := proc(im)
        m,n,d := arrdim(im);
        ans := allocarr[float[8]]([m,n,4]);
        topng0(im,ans,m,n);
        return ans;
    end proc;

    drawpng := proc(arr)
        m,n,d := arrdim(arr);
        if(d=3) then
            return procname(topng(arr));
        end if;
        return ImageTools:-Embed(arr);
    end proc;

#draw a solid circle with radius in pixels at point p
    drawdot := module()
    option object;
    export draw,defsize,defpattern,defcolor,init;
    local ModulePrint,ModuleApply;
        ModulePrint::static := proc()
            return nprintf("draw a solid circle with radius in pixels");
        end proc;
        draw::static := proc(im,p,size:=defsize)
            p1 := im:-xx2ii(p);
            arr := im:-getarr();
            el := [pattern=defpattern,color=defcolor];
            el := prunesubs([op(el),args[4..nargs]]);
            ImageTools:-Draw:-SolidCircle(arr,op(p1),size,op(el));
        end proc;
        ModuleApply::static := draw;
        init::static := proc()
            defsize,defpattern,defcolor := args;
        end proc;
    end module;

    drawdot:-init(4,"solid","black");

#draw a solid circle with radius r at point p
    solidcirc := module()
    option object;
    export draw,defpattern,defcolor,init;
    local ModulePrint,ModuleApply;
        ModulePrint::static := proc()
            return nprintf("draw a solid circle");
        end proc;
        draw::static := proc(im,p,r)
            p1 := im:-xx2ii(p);
            s := im:-getscale();
            arr := im:-getarr();
            el := [pattern=defpattern,color=defcolor];
            el := prunesubs([op(el),args[4..nargs]]);
            ImageTools:-Draw:-SolidCircle(arr,op(p1),r*s,op(el));
        end proc;
        ModuleApply::static := draw;
        init::static := proc()
            defpattern,defcolor := args;
        end proc;
    end module;

    solidcirc:-init("solid","black");

#draw a line from p to q
    drawline := module()
    option object;
    export draw,defthickness,defpattern,defcolor,init;
    local ModulePrint,ModuleApply;
        ModulePrint::static := proc()
            return nprintf("draw a line segment");
        end proc;
        draw::static := proc(im,p,q)
            p1,q1 := im:-xx2ii(p),im:-xx2ii(q);
            arr := im:-getarr();
            el := [thickness=defthickness,pattern=defpattern,color=defcolor];
            el := prunesubs([op(el),args[4..nargs]]);
            ImageTools:-Draw:-Line(arr,op(p1),op(q1),op(el));
        end proc;
        ModuleApply::static := draw;
        init::static := proc()
            defthickness,defpattern,defcolor := args;
        end proc;
    end module;

    drawline:-init(3,"solid","black");

#draw a polygon
    drawpoly := module()
    option object;
    export draw,defthickness,defpattern,defcolor,fillcolor,fillpattern,init;
    local ModulePrint,ModuleApply;
        ModulePrint::static := proc()
            return nprintf("draw a polygon");
        end proc;
        draw::static := proc(im,pl)
            pl1 := map(im:-xx2ii,pl);
            arr := im:-getarr();
            el := [thickness=defthickness,pattern=defpattern,color=defcolor,fill_pattern=fillpattern,fill_color=fillcolor];
            el := prunesubs([op(el),args[3..nargs]]);
            ImageTools:-Draw:-Poly(arr,op(p1),pl1,op(el));
        end proc;
        ModuleApply::static := draw;
        init::static := proc()
            defthickness,defpattern,fillpattern,defcolor,fillcolor := args;
        end proc;
    end module;

    drawpoly:-init(3,"solid","solid","black","white");

    imrng0 := proc()
        md := module()
        option object;
        export getdim,gridsize,draw,getrng,getarr,clear,xx2ii,ii2xx,init,setpixel,`whattype`,scale,getscale;
        local al,bl,ml,d,arr,arr1,width,height;
            ModulePrint::static := proc()
                return nprintf(cat("%dx%d image in [%f..%f,%f..%f]"),ml[1],ml[2],al[1],bl[1],al[2],bl[2]);
            end proc;
            `whattype`::static := proc()
                return 'ImageRange';
            end proc;
            getrng::static := proc()
                return [seq(al[i]..bl[i],i=1..d)];
            end proc;
            getscale::static := proc()
                return scale;
            end proc;
            getdim::static := proc()
                return d;
            end proc;
            gridsize::static := proc()
                return ml;
            end proc;
            setpixel::static := proc(i,j,col)
                i1,j1 := height-j+1,i;
                arr[i1,j1,1],arr[i1,j1,2],arr[i1,j1,3] := col[1],col[2],col[3];
            end proc;
            clear::static := proc()
                ArrayTools:-Fill(1.01,arr);
            end proc;
            getarr::static := proc()
                return arr;
            end proc;
            draw::static := proc()
                topng0(arr,arr1,ml[2],ml[1]);
                return ImageTools:-Embed(arr1);
            end proc;
            init::static := proc()
                al,bl := endrange(args);
                d := nops(al);
                if(d<>2) then
                    error "must be 2-dimensional";
                end if;
                ml,dxl := rastrange(args);
                ii2xx,xx2ii := gridmaps(args);
                arr,arr1 := allocarr[float[8]]([ml[2],ml[1],3],[ml[2],ml[1],4]);
                ArrayTools:-Fill(1.01,arr);
                width,height := ml[1],ml[2];
                scale := max(1/dxl[1],1/dxl[2]);
            end proc;
        end module;
        md:-init(args);
        return md;
    end proc;

    imrng := module()
    option object;
    export defres,tensflag,draw,getim,init;
    local ModulePrint,ModuleApply;
        ModulePrint::static := proc()
            return nprintf("generate spatial images");
        end proc;
        getim::static := proc(rng)
            if(whattype(args[1])='ImageRange') then
                return args[1];
            elif(nargs=1) then
                return imrng0(rng,defres,tensflag);
            elif(nargs=2) then
                return imrng0(args,tensflag);
            else
                return imrng0(args);
            end if;
        end proc;
        ModuleApply::static := getim;
        init::static := proc()
            defres := 1000;
            tensflag := true;
        end proc;
    end module;

    imrng:-init();

#draw a color map of the values of arr up to value c1
    subcolor := proc(arr,c1,col,im)
        m,n := arrdim(arr);
        if(nargs=2) then
            return subcolor(arr,c1,Color([.8,.8,.8]));
        elif(nargs=3) then
            im1 := allocarr[float[8]]([n,m,3]);
            ArrayTools:-Fill(1.1,im1);
            subcolor(arr,c1,col,im1);
            return drawpng(im1);
        end if;
        cmap := colormap(col);
        c0 := min(arr);
        for i from 1 to m do
            for j from 1 to n do
                c := arr[i,j];
                if(c<=c1) then
                    i1,j1 := n-j+1,i;
                    col1 := cmap((c-c0)/(c1-c0));
                    for k from 1 to 3 do
                        im[i1,j1,k] := col1[k];
                    end do;
                end if;
            end do;
        end do;
        return im;
    end proc;

#draw the sublevel sets of a function
    drawsublev := module()
    option object;
    export draw,setcolor,getcolor,cmap,init;
    local ModuleApply,ModulePrint;
        ModulePrint::static := proc()
            return nprintf("draws sublevel sets");
        end proc;
        draw::static := proc(im,f,c1)
            if(whattype(im)<>'ImageRange') then
                error "the image is the first argument now";
            end if;
            #if(whattype(im)<>'ImageRange') then
            #    im1 := imrng(args[3..nargs]);
            #    draw(im1,f,c1);
            #    return im1:-draw();
            #end if;
            rng := im:-getrng();
            gridsize := im:-gridsize();
            A := gridmap(f,rng,gridsize);
            c0 := min(A);
            m,n := op(gridsize);
            for i from 1 to m do
                for j from 1 to n do
                    c := A[i,j];
                    if(c<=c1) then
                        col := cmap((c-c0)/(c1-c0));
                        im:-setpixel(i,j,col);
                    end if;
                end do;
            end do;
            return im;
        end proc;
        ModuleApply::static := draw;
        setcolor::static := proc(col)
            cmap := colormap(col);
        end proc;
        getcolor::static := proc()
            return cmap;
        end proc;
        init::static := proc()
            setcolor(Color([.8,.8,.8]));
            defgrid := 1000;
        end proc;
    end module;

    drawsublev:-init();

#display the objects with background im, on the range given by im
    disprng := proc(cmds,im)
        if(type(procname,indexed)) then
            return disprng([args],op(procname));
        end if;
        arr := im:-getarr();
        rng := im:-getrng();
        return display(cmds,background=arr,view=rng);
    end proc;

    animrng := proc(cmds,im)
        if(type(procname,indexed)) then
            return disprng([args],op(procname));
        end if;
        arr := im:-getarr();
        rng := im:-getrng();
        return display(cmds,insequence=true,background=arr,view=rng);
    end proc;

    drawplot := proc(im,f)
        el := args[3..nargs];
        N := im:-gridsize()[1];
        r1,r2 := op(im:-getrng());
        a,b := op(r1);
        neginf := a-100*(b-a);
        f1 := x->max(f(x),neginf);
        x1 := a+.5/N*(b-a);
        y1 := f1(x1);
        for i from 2 to N do
            x2 := a+(i-.5)/N*(b-a);
            y2 := f1(x2);
            drawline(im,[x1,y1],[x2,y2],el);
            x1,y1 := x2,y2;
        end do;
        return;
    end proc;

    drawvec := proc(im,V,rng)
        el := args[4..nargs];
        N := numelems(V);
        a,b := op(rng);
        x1 := a+.5/N*(b-a);
        for i from 2 to N do
            x2 := a+(i-.5)/N*(b-a);
            drawline(im,[x1,V[i-1]],[x2,V[i]],el);
            x1 := x2;
        end do;
        return;
    end proc;

    drawbelow := proc(im,f)
        el := [args[3..nargs]];
        N := im:-gridsize()[1];
        r1,r2 := op(im:-getrng());
        a1,b1 := op(r1);
        a2,b2 := op(r2);
        neginf := a2-100*(b2-a2);
        f1 := x->max(f(x),neginf);
        eps := (b1-a1)/2;
        pl := [[a1-eps,a2],[a1,f1(a1)]];
        for i from 2 to N do
            x := a1+(i-.5)/N*(b1-a1);
            pl := [op(pl),[x,f1(x)]];
        end do;
        pl := [op(pl),[b1,f1(b1)],[b1+eps,a2]];
        drawpoly(im,pl,op(el));
        return;
    end proc;

end module;
