set role gis;

drop schema if exists hex cascade;
create schema hex;

set search_path to hex,public,gis;

-- return a postgis text polygon
-- the 0,0 hex center is at the origin
-- we scale, rotate, then translate
create or replace function hexgeom(
	xc integer, yc integer, scale float = 1.0, angle float = 0.0, xorigin float = 0.0, yorigin float = 0.0
) returns text language 'plpgsql' as $$
declare
	x float[];
	y float[];
	stride float;
	xoff float;
	yoff float;
	i integer;
	cosa float;
	sina float;
begin
	stride := scale * 2.0 / sqrt(3.0);
	yoff = scale / 2.0;
	x[1] := (2 + 3 * xc ) * stride;
	x[4] := (-2 + 3 * xc) * stride;
	x[2] := (1 + 3 * xc ) * stride;
	x[6] := (1 + 3 * xc ) * stride;
	x[3] := (-1 + 3 * xc ) * stride;
	x[5] := (-1 + 3 * xc ) * stride;

	-- hex coordinates are positive down
	y[1] := (0 - 2 * yc + xc % 2) * yoff;
	y[4] := (0 - 2 * yc + xc % 2) * yoff;
	y[2] := (1 - 2 * yc + xc % 2) * yoff;
	y[3] := (1 - 2 * yc + xc % 2) * yoff;
	y[5] := (-1 - 2 * yc + xc % 2) * yoff;
	y[6] := (-1 - 2 * yc + xc % 2) * yoff;

	if angle != 0.0 then
		cosa := cos(angle);
		sina := sin(angle);

		for i in 1..6 loop
			xoff := cosa * x[i] + sina * y[i];
			y[i] := -sina * x[i] + cosa * y[i];
			x[i] := xoff;
		end loop;
	end if;
	if xorigin != 0.0 or yorigin != 0.0 then
		for i in 1..6 loop
			y[i] := y[i] + yorigin;
			x[i] := x[i] + xorigin;
		end loop;
	end if;

	return format('POLYGON((%s %s, %s %s, %s %s, %s %s, %s %s, %s %s))',
		x[1], y[1], x[2], y[2], x[3], y[3], x[4], y[4], x[5], y[5], x[6], y[6]);
end;
$$ immutable strict;

create or replace function hexcovergrid(
	coverage geometry,
	hexsize numeric
	angle numeric = 0.0,
	full boolean = true,
	srid	integer, -- will default to the srid of the coverage
	origin	geometry(POINT) -- this would be the 0,0 hex center, or perhaps 1,1
) returns table (x integer, y integer, cantor numeric, hex geometry) as $$
declare
	xdim integer,
	ydim integer,
	rcover geometry;
	hex geometry;
	stride double precision;
	xorigin float;
	yorigin float;
	x integer;
	y integer;
	xoff double precision;
begin
	-- the origin will be the center of the 1,1 hex.
	-- if null, then the origin will be set to fit in the
	-- covering box.

	-- counter rotate geometry, fill, rotate grid
	rcover = ST_Rotate(st_transform(coverage,srid), -angle*3.1415926/180.0);
	if srid is null and rcover is not null then
		srid := st_srid(rcover);
	end if;
	
	stride := hexsize * 2.0 / sqrt(3.0);
	-- the center of the 1,1 hex should be one stride right and one half tilesize
	-- down from the upper right of the coverage.  That is, the upper right of
	-- (the bounding box of) the coverage should coincide with the "northwest"
	-- vertex of hex 1,1
	xorigin = st_xmin(rcover) + stride;
	yorigin = st_ymax(rcover) + hexsize / 2.0;

	xdim := ((ST_XMax(rcover) -  ST_XMin(rcover))/hexsize*2/sqrt(3.0))::integer;
	ydim := ((ST_YMax(rcover) -  ST_YMin(rcover))/hexsize)::integer;

	--raise notice 'coverage = % % % %', xmin, ymin, xdim, ydim;
	
        for x in 0 .. xdim loop
                for y in 0 .. ydim loop
			hex := st_geomfromtext(hexgeom(x, ydim-y, hexsize,angle,xorigin,yorigin),srid);
			if not full and not st_intersects(hex, rcover) then
				next;
			end if;
			return query select
		      	 	x, ydim-y,
		       		hex.cantor(ARRAY[x,ydim-y]), 
				hex
			;
                end loop;
        end loop;
END;
$$
language 'plpgsql'
;

create or replace function hexgrid(
	xdim integer,
	ydim integer,
	angle numeric,
	coverage geometry,
	srid	integer,
	origin	geometry(POINT),
	tilesize numeric
) returns table (x integer, y integer, cantor numeric, hex geometry) as $$
declare
	rcover geometry;
	stride double precision;
	x integer;
	y integer;
	xmin double precision;
	ymin double precision;
	xoff double precision;
begin
	-- the origin will be the center of the 1,1 hex.
	-- if null, then the origin will be set to fit in the
	-- covering box.

	-- xdim, and ydim can be null, in which case they are
	-- calculated from the tilesize and the covering box
	
	-- calculate ydim

	-- precedence: ydim, xdim, tilesize

	-- counter rotate geometry, fill, rotate grid
	rcover = ST_Rotate(st_transform(coverage,srid), -angle*3.1415926/180.0);
	if srid is null and rcover is not null then
		srid := st_srid(rcover);
	end if;
	
	xmin := ST_XMin(rcover);
	ymin := ST_YMin(rcover);
	
	stride := tilesize * 2.0 / sqrt(3.0);

	xdim := ((ST_XMax(rcover) -  ST_XMin(rcover))/tilesize*2/sqrt(3.0))::integer;
	ydim := ((ST_YMax(rcover) -  ST_YMin(rcover))/tilesize)::integer;

	--raise notice 'coverage = % % % %', xmin, ymin, xdim, ydim;
	
        for x in 0 .. xdim loop
                for y in 0 .. ydim loop
			return query select
		      	 	x, ydim-y,
		       		hex.cantor(ARRAY[x,ydim-y]), 
				st_geomfromtext(hexgeom(x, ydim-y, tilesize,angle),srid)
			;
                end loop;
        end loop;
END;
$$
language 'plpgsql'
;

-- create some hexagon ops

create function ijk(xy integer[]) returns integer[] language 'plpgsql' as $$
declare
	ijk integer[];
begin
	if array_length(xy,1) = 3 then
		return xy;
	end if;

	ijk[1] := xy[1];
	ijk[2] := -xy[2];

	if xy[1] < 0 then
		ijk[2] := ijk[2] + (-xy[1]+1) / 2;
	else
		ijk[2] := ijk[2] - xy[1]/2;
	end if;

	ijk[3] := -ijk[1] - ijk[2];

	return ijk;
end
$$ immutable strict;

create function ijk(x integer, y integer) returns integer[] language 'sql' as $$
	select hex.ijk(ARRAY[x,y]);
$$ immutable strict;

-- invert the cantor paring
create function xy(hex numeric) returns integer[] language 'plpgsql' as $$
declare
	xy integer[];
	w integer;
	t integer;
	py integer;
begin
	if hex >= 0 then
		hex := hex-1;
		w := floor((sqrt(8.0 * hex + 1.0) - 1.0)/2.0)::integer;
		t := (w * w + w)/2;
		py := hex - t;

		xy = ARRAY[w-py,py];
	-- handle the negative cases
	else 
		hex := -hex;
		hex := hex-1;
		w := floor((sqrt(8.0 * hex + 1.0) - 1.0)/2.0)::integer;
		t := (w * w + w)/2;
		py := hex - t;

		xy = ARRAY[w-py,py];

		if xy[1] % 2 = 1 then
			xy[1] = -(xy[1]+1)/2;
		else
			xy[1] = xy[1]/2;
		end if;
		if xy[2] % 2 = 1 then
			xy[2] = -(xy[2]+1)/2;
		else
			xy[2] = xy[2]/2;
		end if;
	end if;

	return xy;
end;
$$ immutable strict;

create function ijk(hex numeric) returns integer[] language 'sql' as $$
	select hex.ijk(hex.xy(hex));
$$ immutable strict;

create function xy(ijk integer[]) returns integer[] language 'plpgsql' as $$
declare
	xy integer[];
begin
	-- make sure this is really an ijk
	-- just no-op if it's already an
	-- xy, so this is sort of a coercion function
	if array_length(ijk,1) = 2 then
		return ijk;
	end if;
	xy[1] := ijk[1];
	xy[2] := -ijk[2];

	if ijk[1] < 0 then
		xy[2] := xy[2] + (-ijk[1] + 1) / 2;
	else
		xy[2] := xy[2] - ijk[1]/2;
	end if;
	return xy;
end;
$$;

CREATE OR REPLACE FUNCTION vec_add(arr1 integer[], arr2 integer[])
RETURNS integer[] AS
$$
SELECT array_agg(result)
FROM (SELECT tuple.val1 + tuple.val2 AS result
      FROM (SELECT UNNEST($1) AS val1
                   ,UNNEST($2) AS val2
                   ,generate_subscripts($1, 1) AS ix) tuple
      ORDER BY ix) inn;
$$ LANGUAGE SQL IMMUTABLE STRICT;

create function adjhex(dir integer, hex integer[]) returns integer[] language 'plpgsql' as $$
declare
	adj integer[];
begin
	if array_length(hex,1) = 2 then
		hex = hex.ijk(hex);
	end if;

	dir := dir % 6;
	if dir = 2 then
		hex := hex.vec_add(hex, ARRAY[-1, 1, 0]);
	elsif dir = 1 then
		hex := hex.vec_add(hex, ARRAY[ 0, 1,-1]);
	elsif dir = 0 then
		hex := hex.vec_add(hex, ARRAY[ 1, 0,-1]);
	elsif dir = 5 then
		hex := hex.vec_add(hex, ARRAY[ 1,-1, 0]);
	elsif dir = 4 then
		hex := hex.vec_add(hex, ARRAY[ 0,-1, 1]);
	elsif dir = 3 then
		hex := hex.vec_add(hex, ARRAY[-1, 0, 1]);
	else
		-- error
	end if;
	return hex;
end;
$$ immutable strict;

create function adjhex(dir integer, hex numeric) returns integer[] language 'sql' as $$
	select hex.adjhex(dir, hex.ijk(hex));
$$ immutable strict;

-- TODO make this work for any, or inline it
create function natcantor(hex integer[]) returns numeric language 'sql' as $$
	select cast((hex[1] + hex[2]) * (hex[1] + hex[2] + 1) / 2 + hex[2] + 1 as numeric);
$$ immutable strict;

create function cantor(hex integer[]) returns numeric language 'plpgsql' as $$
begin
	-- convert to xy
	if array_length(hex,1) = 3 then
		hex = hex.xy(hex);
	end if;

	if hex[1] < 0 or hex[2] < 0 then
		hex[1] = hex[1] * 2;
		hex[2] = hex[2] * 2;
		if hex[1] < 0 then
			hex[1] = -hex[1] - 1;
		end if;
		if hex[2] < 0 then
			hex[2] = -hex[2] - 1;
		end if;
		return -cast((hex[1] + hex[2]) * (hex[1] + hex[2] + 1) / 2 + hex[2] + 1 as numeric);
	end if;

	return cast((hex[1] + hex[2]) * (hex[1] + hex[2] + 1) / 2 + hex[2] + 1 as numeric);
end;
$$ immutable strict;

create function adjacent(hex numeric) returns setof numeric language 'sql' as $$
select hex.cantor(hex.adjhex(dir, hex)) from generate_series(0,5) as dir;
$$ immutable strict;

create function adjacent(dir integer, hex numeric) returns numeric language 'sql' as $$
	select hex.cantor(hex.adjhex(dir, hex.ijk(hex)));
$$ immutable strict;

create function within(hex numeric, range integer = 1, include boolean = false) returns setof numeric language 'plpgsql' as $$
declare
	c integer[]; -- cantor coordinates
	base integer[];
begin
	if range <= 0 then
		return;
	end if;
	if include then
		return next hex;
	end if;

	-- we could in theory eliminate the inner loop and
	-- do a return query
	base = hex.ijk(hex);
	for r in 1 .. range loop
		c = base;
		c[1] := c[1] + r;
		c[3] := -c[1] - c[2]; -- keep the invariant
		hex := hex.cantor(c);
		
		for h in 0 .. r * 6 - 1 loop
			return next hex;
			hex = hex.adjacent(h/r+2,hex);
		end loop;
	end loop;
	return;
end;
$$ immutable strict;

create function within_array(hex numeric, range integer = 1, include boolean = false) returns numeric[] language 'sql' as $$
select array_agg(c) from hex.within(hex, range, include) c;
$$ immutable strict;

create function atrange(hex numeric, range integer = 1) returns setof numeric language 'plpgsql' as $$
declare
	c integer[]; -- cantor coordinates
begin
	if range = 0 then
		return next hex;
	end if;
	if range <= 0 then
		return;
	end if;

	c = ijk(hex);
	c[1] := c[1] + range;
	c[3] := -c[1] - c[2]; -- keep the invariant
	hex := hex.cantor(c);
	
	for q in 0 .. range * 6 - 1 loop
		return next hex;
		hex = hex.adjacent(q/range+2,hex);
	end loop;
	return;
end;
$$ immutable strict;

select
xy(ARRAY[1,2]),
ijk(xy(ARRAY[1,2])),
xy(ijk(xy(ARRAY[1,2])))
;

select xy(ARRAY[2,2]), dir, xy(adjhex(dir, ijk(ARRAY[2,2]))) from generate_series(0,5) as dir;
select
xy(adjhex(dir, ijk(ARRAY[1,2]))),
cantor(xy(adjhex(dir, ijk(ARRAY[1,2])))),
xy(cantor(xy(adjhex(dir, ijk(ARRAY[1,2])))))
from generate_series(0,5) as dir;

select cantor(ARRAY[4,3]);
--select * from within(cantor(ARRAY[4,4]), 2);
select h, xy(h) from within(cantor(ARRAY[4,3]), 2) as h;