#include "globals.h"
#include "flags.h"
#include "usb/usb.h"
#include "bitmap.h"
#include "fluxmap.h"
#include "decoders/fluxmapreader.h"
#include "writer.h"
#include "protocol.h"
#include "proto.h"
#include "fluxsink/fluxsink.h"
#include "fmt/format.h"
#include "dep/agg/include/agg2d.h"
#include "dep/stb/stb_image_write.h"
#include <fstream>

static FlagGroup flags;

static StringFlag destFlux(
	{ "--dest", "-d" },
	"'drive:' flux destination to analyse",
	"",
	[](const auto& value)
	{
		FluxSink::updateConfigForFilename(config.mutable_flux_sink(), value);
	});

static IntFlag destCylinder(
	{ "--cylinder", "-c" },
	"cylinder to write to",
	0);

static IntFlag destHead(
	{ "--head", "-h" },
	"head to write to",
	0);

static DoubleFlag minInterval(
	{ "--min-interval-us" },
	"Minimum pulse interval",
	2.0);

static DoubleFlag maxInterval(
	{ "--max-interval-us" },
	"Maximum pulse interval",
	10.0);

static DoubleFlag intervalStep(
	{ "--interval-step-us" },
	"Interval step, approximately",
	0.2);

static StringFlag writeCsv(
	{ "--write-csv" },
	"Write detailed CSV data",
	"");

static StringFlag writeImg(
	{ "--write-img" },
	"Draw a graph of the response data",
	"analysis.png");

static IntFlag imgWidth(
	{ "--width" },
	"Width of output graph",
	800);

static IntFlag imgHeight(
	{ "--height" },
	"Height of output graph",
	600);

/* This is the Turbo colourmap.
 * https://ai.googleblog.com/2019/08/turbo-improved-rainbow-colormap-for.html
 */
static const uint8_t turbo_srgb_bytes[256][3] = {
	 {48,18,59}, {50,21,67}, {51,24,74}, {52,27,81}, {53,30,88}, {54,33,95},
	 {55,36,102}, {56,39,109}, {57,42,115}, {58,45,121}, {59,47,128},
	 {60,50,134}, {61,53,139}, {62,56,145}, {63,59,151}, {63,62,156},
	 {64,64,162}, {65,67,167}, {65,70,172}, {66,73,177}, {66,75,181},
	 {67,78,186}, {68,81,191}, {68,84,195}, {68,86,199}, {69,89,203},
	 {69,92,207}, {69,94,211}, {70,97,214}, {70,100,218}, {70,102,221},
	 {70,105,224}, {70,107,227}, {71,110,230}, {71,113,233}, {71,115,235},
	 {71,118,238}, {71,120,240}, {71,123,242}, {70,125,244}, {70,128,246},
	 {70,130,248}, {70,133,250}, {70,135,251}, {69,138,252}, {69,140,253},
	 {68,143,254}, {67,145,254}, {66,148,255}, {65,150,255}, {64,153,255},
	 {62,155,254}, {61,158,254}, {59,160,253}, {58,163,252}, {56,165,251},
	 {55,168,250}, {53,171,248}, {51,173,247}, {49,175,245}, {47,178,244},
	 {46,180,242}, {44,183,240}, {42,185,238}, {40,188,235}, {39,190,233},
	 {37,192,231}, {35,195,228}, {34,197,226}, {32,199,223}, {31,201,221},
	 {30,203,218}, {28,205,216}, {27,208,213}, {26,210,210}, {26,212,208},
	 {25,213,205}, {24,215,202}, {24,217,200}, {24,219,197}, {24,221,194},
	 {24,222,192}, {24,224,189}, {25,226,187}, {25,227,185}, {26,228,182},
	 {28,230,180}, {29,231,178}, {31,233,175}, {32,234,172}, {34,235,170},
	 {37,236,167}, {39,238,164}, {42,239,161}, {44,240,158}, {47,241,155},
	 {50,242,152}, {53,243,148}, {56,244,145}, {60,245,142}, {63,246,138},
	 {67,247,135}, {70,248,132}, {74,248,128}, {78,249,125}, {82,250,122},
	 {85,250,118}, {89,251,115}, {93,252,111}, {97,252,108}, {101,253,105},
	 {105,253,102}, {109,254,98}, {113,254,95}, {117,254,92}, {121,254,89},
	 {125,255,86}, {128,255,83}, {132,255,81}, {136,255,78}, {139,255,75},
	 {143,255,73}, {146,255,71}, {150,254,68}, {153,254,66}, {156,254,64},
	 {159,253,63}, {161,253,61}, {164,252,60}, {167,252,58}, {169,251,57},
	 {172,251,56}, {175,250,55}, {177,249,54}, {180,248,54}, {183,247,53},
	 {185,246,53}, {188,245,52}, {190,244,52}, {193,243,52}, {195,241,52},
	 {198,240,52}, {200,239,52}, {203,237,52}, {205,236,52}, {208,234,52},
	 {210,233,53}, {212,231,53}, {215,229,53}, {217,228,54}, {219,226,54},
	 {221,224,55}, {223,223,55}, {225,221,55}, {227,219,56}, {229,217,56},
	 {231,215,57}, {233,213,57}, {235,211,57}, {236,209,58}, {238,207,58},
	 {239,205,58}, {241,203,58}, {242,201,58}, {244,199,58}, {245,197,58},
	 {246,195,58}, {247,193,58}, {248,190,57}, {249,188,57}, {250,186,57},
	 {251,184,56}, {251,182,55}, {252,179,54}, {252,177,54}, {253,174,53},
	 {253,172,52}, {254,169,51}, {254,167,50}, {254,164,49}, {254,161,48},
	 {254,158,47}, {254,155,45}, {254,153,44}, {254,150,43}, {254,147,42},
	 {254,144,41}, {253,141,39}, {253,138,38}, {252,135,37}, {252,132,35},
	 {251,129,34}, {251,126,33}, {250,123,31}, {249,120,30}, {249,117,29},
	 {248,114,28}, {247,111,26}, {246,108,25}, {245,105,24}, {244,102,23},
	 {243,99,21}, {242,96,20}, {241,93,19}, {240,91,18}, {239,88,17},
	 {237,85,16}, {236,83,15}, {235,80,14}, {234,78,13}, {232,75,12},
	 {231,73,12}, {229,71,11}, {228,69,10}, {226,67,10}, {225,65,9},
	 {223,63,8}, {221,61,8}, {220,59,7}, {218,57,7}, {216,55,6}, {214,53,6},
	 {212,51,5}, {210,49,5}, {208,47,5}, {206,45,4}, {204,43,4}, {202,42,4},
	 {200,40,3}, {197,38,3}, {195,37,3}, {193,35,2}, {190,33,2}, {188,32,2},
	 {185,30,2}, {183,29,2}, {180,27,1}, {178,26,1}, {175,24,1}, {172,23,1},
	 {169,22,1}, {167,20,1}, {164,19,1}, {161,18,1}, {158,16,1}, {155,15,1},
	 {152,14,1}, {149,13,1}, {146,11,1}, {142,10,1}, {139,9,2}, {136,8,2},
	 {133,7,2}, {129,6,2}, {126,5,2}, {122,4,3}};

static void palette(double value, agg::srgba8* pixel)
{
	int index = std::min((int)(value * 256.0), 255);
	pixel->r = turbo_srgb_bytes[index][0];
	pixel->g = turbo_srgb_bytes[index][1];
	pixel->b = turbo_srgb_bytes[index][2];
	pixel->a = 255;
}

static void do_in_steps(double c1, double c2, double lo, double hi,
		double step, std::function<void(double c, double v)> func)
{
	double scale = (c2 - c1) / (hi - lo);
	double v = lo;
	while (v <= hi + step/10.0)
	{
		double c = c1 + scale*(v - lo);
		func(c, v);
		v += step;
	}
}

static void draw_y_axis(Agg2D& painter, double x, double y1, double y2,
		double lo, double hi, double step, const char* format)
{
	painter.noFill();
	painter.lineColor(0, 0, 0);
	painter.line(x, y1, x, y2);
	painter.textSize(10.0);
	painter.textAlignment(Agg2D::AlignRight);

	do_in_steps(y1, y2, lo, hi, step, [&](double y, double v)
	{
		painter.line(x, y, x-5, y);
		painter.text(x-8, y+5.0, fmt::format(format, v));
	});
}

static void draw_x_axis(Agg2D& painter, double x1, double x2, double y,
		double lo, double hi, double step, const char* format)
{
	painter.noFill();
	painter.lineColor(0, 0, 0);
	painter.line(x1, y, x2, y);
	painter.textSize(10.0);
	painter.textAlignment(Agg2D::AlignCenter);

	do_in_steps(x1, x2, lo, hi, step, [&](double x, double v)
	{
		painter.line(x, y, x, y+5);
		painter.text(x, y+18, fmt::format(format, v));
	});
}

static void draw_y_graticules(Agg2D& painter, double x1, double y1, double x2, double y2,
		double lo, double hi, double step)
{
	painter.noFill();
	painter.lineColor(0, 0, 0, 128);

	do_in_steps(y1, y2, lo, hi, step, [&](double y, double v)
	{
		painter.line(x1, y, x2, y);
	});
}

static void draw_x_graticules(Agg2D& painter, double x1, double y1, double x2, double y2,
		double lo, double hi, double step)
{
	painter.noFill();
	painter.lineColor(0, 0, 0, 128);

	do_in_steps(x1, x2, lo, hi, step, [&](double x, double v)
	{
		painter.line(x, y1, x, y2);
	});
}

int mainAnalyseDriveResponse(int argc, const char* argv[])
{
    flags.parseFlagsWithConfigFiles(argc, argv, {});

	if (!config.flux_sink().has_drive())
		Error() << "this only makes sense with a real disk drive";

    usbSetDrive(config.flux_sink().drive().drive(),
		config.flux_sink().drive().high_density(),
		config.flux_sink().drive().index_mode());
	usbSeek(destCylinder);

	std::cout << "Measuring rotational speed...\n";
    nanoseconds_t period = usbGetRotationalPeriod(0);
	if (period == 0)
		Error() << "Unable to measure rotational speed (try fluxengine rpm).";

	std::ofstream csv;
	if (writeCsv.get() != "")
		csv.open(writeCsv);

	int numRows = (maxInterval - minInterval) / intervalStep;
	const int numColumns = 512;
	std::vector<std::vector<double>> frequencies(numRows, std::vector<double>(numColumns, 0.0));

	double interval;
	for (int row=0; row<numRows; row++)
	{
		interval = minInterval + (double)row * intervalStep;

		unsigned ticks = (unsigned) (interval * TICKS_PER_US);
		std::cout << fmt::format("Interval {:.2f}: ", ticks * US_PER_TICK);
		std::cout << std::flush;

		if (interval >= 2.0)
		{
			/* Write the test pattern. */

			Fluxmap outFluxmap;
			while (outFluxmap.duration() < period)
			{
				outFluxmap.appendInterval(ticks);
				outFluxmap.appendPulse();
			}
			usbWrite(destHead, outFluxmap.rawBytes(), 0);

			/* Read the test pattern in again. */

			Fluxmap inFluxmap;
			inFluxmap.appendBytes(usbRead(destHead, true, period, 0));

			/* Compute histogram. */

			FluxmapReader fmr(inFluxmap);
			fmr.seek((double)period*0.1); /* skip first 10% and last 10% as contains junk */
			fmr.skipToEvent(F_BIT_PULSE);
			while (fmr.tell().ns() < ((double)period*0.9))
			{
				unsigned ticks;
				fmr.findEvent(F_BIT_PULSE, ticks);
				if (ticks < numColumns)
					frequencies[row][ticks]++;
			}
		}

		/* Compute mean and normalise. */

		double sum = 0.0;
		double prod = 0.0;
		double max = 0.0;
		for (int i=0; i<numColumns; i++)
		{
			sum += frequencies[row][i];
			prod += i * frequencies[row][i];
			max = std::max(max, frequencies[row][i]);
		}
		if (max != 0.0)
			for (int i=0; i<numColumns; i++)
				frequencies[row][i] /= max;

		if (sum == 0)
			std::cout << "failed\n";
		else
		{
			double mean = prod / sum;
			std::cout << fmt::format("{:.4f}\n", mean/TICKS_PER_US);

		}

		if (writeCsv.get() != "")
		{
			csv << interval;
			for (double d : frequencies[row])
				csv << "," << d;
			csv << '\n';
		}
	}

	Bitmap bitmap(writeImg, imgWidth, imgHeight);
	if (!bitmap.filename.empty())
	{
		Agg2D& painter = bitmap.painter();
		painter.clearAll(0xdd, 0xdd, 0xdd);

		const double MARGIN = 30;
		agg::rect_d drawableBounds = {
			MARGIN*1.5, MARGIN,
			bitmap.width - MARGIN, bitmap.height - MARGIN
		};
		agg::rect_d colourbarBounds = {
			drawableBounds.x2 - MARGIN, drawableBounds.y1,
			drawableBounds.x2, drawableBounds.y2
		};
		agg::rect_d graphBounds = {
			drawableBounds.x1, drawableBounds.y1,
			colourbarBounds.x1 - MARGIN*2, drawableBounds.y2
		};
		double blockWidth = (graphBounds.x2 - graphBounds.x1) / numColumns;
		double blockHeight = (graphBounds.y2 - graphBounds.y1) / numRows;
		painter.imageFilter(Agg2D::NoFilter);
		painter.imageResample(Agg2D::NoResample);
		painter.imageBlendMode(Agg2D::BlendDst);

		/* Create the off-screen buffer which the actual bitmap goes into, and draw it. */

		{
			const int width = numRows; /* input interval on X axis */
			const int height = numColumns; /* response spread on Y axis */
			std::vector<agg::srgba8> rbufdata(height*width);
			for (int y=0; y<height; y++)
				for (int x=0; x<width; x++)
					palette(frequencies[x][y], &rbufdata[x + y*width]);

			Agg2D::Image image((uint8_t*)&rbufdata[0], width, height, width*sizeof(agg::srgba8));
			painter.transformImage(image, graphBounds.x1, graphBounds.y2, graphBounds.x2, graphBounds.y1);
		}

		/* Likewise for the colour bar. */

		{
			const int height = graphBounds.y2 - graphBounds.y1;
			std::vector<agg::srgba8> rbufdata(height);
			for (int y=0; y<height; y++)
				palette((double)y / height, &rbufdata[y]);

			Agg2D::Image image((uint8_t*)&rbufdata[0], 1, height, sizeof(agg::srgba8));
			painter.transformImage(image, colourbarBounds.x1, colourbarBounds.y2, colourbarBounds.x2, colourbarBounds.y1);
		}

		draw_y_axis(painter, colourbarBounds.x1-5, colourbarBounds.y2, colourbarBounds.y1, 0.0, 1.0, 0.1, "{:.1f}");
		draw_y_axis(painter, graphBounds.x1-5, graphBounds.y2, graphBounds.y1, 0.0, 512.0/TICKS_PER_US, 5.0, "{:.0f}");
		draw_y_graticules(painter, graphBounds.x1, graphBounds.y2, graphBounds.x2, graphBounds.y1, 0.0, 512.0/TICKS_PER_US, 5.0);
		draw_x_axis(painter, graphBounds.x1, graphBounds.x2, graphBounds.y2+5, minInterval, maxInterval, 5.0, "{:.0f}");
		draw_x_graticules(painter, graphBounds.x1, graphBounds.y1, graphBounds.x2, graphBounds.y2, minInterval, maxInterval, 5.0);

		painter.noFill();
		painter.lineColor(0, 0, 0);
		painter.rectangle(graphBounds.x1, graphBounds.y1, graphBounds.x2, graphBounds.y2);
		painter.rectangle(colourbarBounds.x1, drawableBounds.y1, drawableBounds.x2, drawableBounds.y2);
		bitmap.save();
	}

    return 0;
}