#include "dynamic_tree.h"
#include <algorithm>
#include <stdexcept>
#include <functional>
#include <cmath>
namespace noita::cpu {
namespace {
ProxyAabb unite(ProxyAabb a,ProxyAabb b){return {{std::min(a.lower.x,b.lower.x),std::min(a.lower.y,b.lower.y)},
    {std::max(a.upper.x,b.upper.x),std::max(a.upper.y,b.upper.y)}};}
double perimeter(ProxyAabb a){return 2*((a.upper.x-a.lower.x)+(a.upper.y-a.lower.y));}
bool same(ProxyAabb a,ProxyAabb b){return a.lower.x==b.lower.x&&a.lower.y==b.lower.y&&a.upper.x==b.upper.x&&a.upper.y==b.upper.y;}
}
bool aabb_overlap(ProxyAabb a,ProxyAabb b){return !(a.lower.x>b.upper.x||a.lower.y>b.upper.y||b.lower.x>a.upper.x||b.lower.y>a.upper.y);}
DynamicTree::DynamicTree():nodes_(16){for(int id=15;id>=0;--id)free_.push_back(id);}
int DynamicTree::allocate(){if(free_.empty()){int old=int(nodes_.size());nodes_.resize(old*2);
        for(int id=int(nodes_.size())-1;id>=old;--id)free_.push_back(id);}
    int id=free_.back();free_.pop_back();nodes_[id]=Node{};nodes_[id].height=0;return id;}
void DynamicTree::release(int id){nodes_[id]=Node{};free_.push_back(id);}
void DynamicTree::refresh(int id){auto& n=nodes_[id];n.box=unite(nodes_[n.left].box,nodes_[n.right].box);
    n.height=1+std::max(nodes_[n.left].height,nodes_[n.right].height);}
int DynamicTree::balance(int a){if(nodes_[a].left==-1||nodes_[a].height<2)return a;
    int b=nodes_[a].left,c=nodes_[a].right,diff=nodes_[c].height-nodes_[b].height;
    if(diff>1){int f=nodes_[c].left,g=nodes_[c].right,p=nodes_[a].parent;
        nodes_[c].left=a;nodes_[c].parent=p;nodes_[a].parent=c;
        if(p==-1)root_=c;else if(nodes_[p].left==a)nodes_[p].left=c;else nodes_[p].right=c;
        if(nodes_[f].height>nodes_[g].height){nodes_[c].right=f;nodes_[a].right=g;nodes_[g].parent=a;}
        else{nodes_[c].right=g;nodes_[a].right=f;nodes_[f].parent=a;}
        refresh(a);refresh(c);return c;}
    if(diff<-1){int d=nodes_[b].left,e=nodes_[b].right,p=nodes_[a].parent;
        nodes_[b].left=a;nodes_[b].parent=p;nodes_[a].parent=b;
        if(p==-1)root_=b;else if(nodes_[p].left==a)nodes_[p].left=b;else nodes_[p].right=b;
        if(nodes_[d].height>nodes_[e].height){nodes_[b].right=d;nodes_[a].left=e;nodes_[e].parent=a;}
        else{nodes_[b].right=e;nodes_[a].left=d;nodes_[d].parent=a;}
        refresh(a);refresh(b);return b;}
    return a;}
void DynamicTree::insert(int leaf){if(root_==-1){root_=leaf;nodes_[leaf].parent=-1;return;}
    int sibling=root_;auto box=nodes_[leaf].box;
    while(nodes_[sibling].left!=-1){int a=nodes_[sibling].left,b=nodes_[sibling].right;
        double area=perimeter(nodes_[sibling].box),combined=perimeter(unite(nodes_[sibling].box,box));
        double direct=2*combined,inherit=2*(combined-area);
        auto cost=[&](int i){double value=perimeter(unite(nodes_[i].box,box));
            return value+inherit-(nodes_[i].left==-1?0:perimeter(nodes_[i].box));};
        double ca=cost(a),cb=cost(b);if(direct<ca&&direct<cb)break;
        sibling=ca<cb?a:b; // Equal costs choose child2 in original.
    }
    int old=nodes_[sibling].parent,parent=allocate();nodes_[parent].parent=old;
    nodes_[parent].left=sibling;nodes_[parent].right=leaf;nodes_[sibling].parent=parent;nodes_[leaf].parent=parent;
    refresh(parent);if(old==-1)root_=parent;else if(nodes_[old].left==sibling)nodes_[old].left=parent;else nodes_[old].right=parent;
    for(int i=parent;i!=-1;){i=balance(i);refresh(i);i=nodes_[i].parent;}}
void DynamicTree::detach(int leaf){if(root_==leaf){root_=-1;return;}
    int p=nodes_[leaf].parent,g=nodes_[p].parent,s=nodes_[p].left==leaf?nodes_[p].right:nodes_[p].left;
    if(g==-1){root_=s;nodes_[s].parent=-1;release(p);}
    else{if(nodes_[g].left==p)nodes_[g].left=s;else nodes_[g].right=s;nodes_[s].parent=g;release(p);
        for(int i=g;i!=-1;){i=balance(i);refresh(i);i=nodes_[i].parent;}}
    nodes_[leaf].parent=-1;}
int DynamicTree::create(ProxyAabb box,int user){int id=allocate();nodes_[id].box=box;nodes_[id].user=user==-1?id:user;insert(id);return id;}
void DynamicTree::destroy(int id){if(nodes_.at(id).height!=0)throw std::invalid_argument("not a live tree leaf");detach(id);release(id);}
void DynamicTree::move(int id,ProxyAabb box){if(nodes_.at(id).height!=0)throw std::invalid_argument("not a live tree leaf");detach(id);nodes_[id].box=box;insert(id);}
std::vector<int> DynamicTree::query(ProxyAabb box) const {std::vector<int> out,stack{root_};
    while(!stack.empty()){int i=stack.back();stack.pop_back();if(i==-1||!aabb_overlap(box,nodes_[i].box))continue;
        const auto& n=nodes_[i];if(n.left==-1)out.push_back(n.user);else{stack.push_back(n.left);stack.push_back(n.right);}}
    return out;}
int DynamicTree::height() const{return root_==-1?0:nodes_[root_].height;}
void DynamicTree::raycast(Vec2d start,Vec2d end,double fraction,const std::function<double(int,double)>& callback) const {
    if(!callback)throw std::invalid_argument("ray callback required");Vec2d delta{end.x-start.x,end.y-start.y},unit=delta;
    double length=std::hypot(unit.x,unit.y);if(length>=1.1920928955078125e-7){unit.x/=length;unit.y/=length;}
    Vec2d perpendicular{-unit.y,unit.x};auto segment=[&]{Vec2d p{start.x+fraction*delta.x,start.y+fraction*delta.y};
        return ProxyAabb{{std::min(start.x,p.x),std::min(start.y,p.y)},{std::max(start.x,p.x),std::max(start.y,p.y)}};};
    auto box=segment();std::vector<int> stack{root_};while(!stack.empty()){int i=stack.back();stack.pop_back();if(i==-1||!aabb_overlap(box,nodes_[i].box))continue;
        const auto& n=nodes_[i];Vec2d center{(n.box.lower.x+n.box.upper.x)*.5,(n.box.lower.y+n.box.upper.y)*.5},
            extent{(n.box.upper.x-n.box.lower.x)*.5,(n.box.upper.y-n.box.lower.y)*.5};
        double separation=std::abs(perpendicular.x*(start.x-center.x)+perpendicular.y*(start.y-center.y))
            -std::abs(perpendicular.x)*extent.x-std::abs(perpendicular.y)*extent.y;
        if(separation>0)continue;if(n.left!=-1){stack.push_back(n.left);stack.push_back(n.right);continue;}
        double value=callback(n.user,fraction);if(value==0)return;if(value>0){fraction=value;box=segment();}}
}
bool DynamicTree::valid() const {std::vector<bool> seen(nodes_.size());bool ok=true;
    std::function<int(int,int)> visit=[&](int i,int parent){if(i<0||i>=int(nodes_.size())||seen[i]){ok=false;return -1;}
        seen[i]=true;const auto& n=nodes_[i];if(n.parent!=parent||n.height<0)ok=false;
        if(n.left==-1){if(n.height!=0||n.right!=-1)ok=false;return 0;}
        int a=visit(n.left,i),b=visit(n.right,i);if(a<0||b<0)return -1;
        if(n.height!=1+std::max(a,b)||!same(n.box,unite(nodes_[n.left].box,nodes_[n.right].box)))ok=false;
        return n.height;};
    if(root_!=-1)visit(root_,-1);
    for(int i:free_){if(i<0||i>=int(seen.size())||seen[i]||nodes_[i].height!=-1)ok=false;else seen[i]=true;}
    for(bool s:seen)if(!s)ok=false;return ok;}
}