#include "catalog.h"

#include <math.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <omp.h>

#define PI 3.14159265358979323846

enum {
    TEST_GRID_LINE_DEG = 10,
    TEST_GRID_SAMPLE_DEG = 2,
    TEST_GRID_RED_TEMPERATURE_K = 3000,
    TEST_GRID_BLUE_TEMPERATURE_K = 12000,
};

/*
 * With the renderer's 380--780 nm CIE integration and linear-sRGB luminance,
 * B(3000 K) / B(12000 K) = 0.00141095580387.  Red stars retain unit scale.
 */
#define TEST_GRID_BLUE_AMPLITUDE 0.00141095580387

static int test_grid_temperature_K(int longitude_deg, int latitude_deg)
{
    /* Boundaries belong to the octant immediately east/north of them. */
    const int longitude_sector = longitude_deg / 90;
    const int hemisphere = latitude_deg < 0 ? 0 : 1;
    /* Add the hemisphere bit to flip the color across the equator. */
    const int octant = hemisphere + longitude_sector;
    return octant % 2 == 0 ? TEST_GRID_RED_TEMPERATURE_K
                           : TEST_GRID_BLUE_TEMPERATURE_K;
}

static double test_grid_amplitude(int temperature_K)
{
    return temperature_K == TEST_GRID_BLUE_TEMPERATURE_K
               ? TEST_GRID_BLUE_AMPLITUDE
               : 1.0;
}

int catalog_write_octant_grid(const char *path)
{
    FILE *file = fopen(path, "w");
    if (file == NULL) return -1;
    fputs("longitude_deg,latitude_deg,temperature_K,amplitude\n", file);
    for (int latitude = -90; latitude <= 90; latitude += TEST_GRID_SAMPLE_DEG) {
        for (int longitude = 0; longitude < 360;
             longitude += TEST_GRID_SAMPLE_DEG) {
            const int on_longitude_line = longitude % TEST_GRID_LINE_DEG == 0;
            const int on_latitude_line = latitude % TEST_GRID_LINE_DEG == 0;
            if ((!on_longitude_line && !on_latitude_line) ||
                ((latitude == -90 || latitude == 90) && longitude != 0))
                continue;
            const int temperature_K =
                test_grid_temperature_K(longitude, latitude);
            fprintf(file, "%d,%d,%d,%.12g\n", longitude, latitude,
                    temperature_K, test_grid_amplitude(temperature_K));
        }
    }
    return fclose(file) == 0 ? 0 : -1;
}

int catalog_load_csv(StarCatalog *catalog, const char *path)
{
    FILE *file = fopen(path, "r");
    char line[256];
    size_t capacity = 0;
    if (catalog == NULL) {
        if (file != NULL) fclose(file);
        return -1;
    }
    *catalog = (StarCatalog){0};
    if (file == NULL || fgets(line, sizeof line, file) == NULL) goto fail;

    while (fgets(line, sizeof line, file) != NULL) {
        double longitude, latitude, temperature, amplitude;
        if (sscanf(line, "%lf,%lf,%lf,%lf", &longitude, &latitude,
                   &temperature, &amplitude) != 4)
            goto fail;
        if (catalog->count == capacity) {
            size_t next = capacity == 0 ? 256 : capacity * 2;
            Star *stars = realloc(catalog->stars, next * sizeof *stars);
            if (stars == NULL) goto fail;
            catalog->stars = stars;
            capacity = next;
        }
        const double lon = longitude * PI / 180.0;
        const double lat = latitude * PI / 180.0;
        const double cos_lat = cos(lat);
        Star *star = &catalog->stars[catalog->count++];
        /* Standard right-handed ICRS Cartesian axes: X is (RA, Dec) =
         * (0, 0), Y is (90, 0), and Z is the north celestial pole. */
        star->direction[0] = cos_lat * cos(lon);
        star->direction[1] = cos_lat * sin(lon);
        star->direction[2] = sin(lat);
        star->temperature_K = temperature;
        star->amplitude = amplitude;
    }
    fclose(file);
    return 0;
fail:
    if (file != NULL) fclose(file);
    catalog_destroy(catalog);
    return -1;
}

static size_t tile_index(int ra_index, int dec_index)
{
    return (size_t)dec_index * CATALOG_ALL_SKY_RA_TILES + ra_index;
}

static int load_tile_file(const StarCatalog *catalog, int ra_index,
                          int dec_index, Star **stars, size_t *count)
{
    char path[PATH_MAX];
    int written;
    if (catalog == NULL || stars == NULL || count == NULL)
        return -1;
    *stars = NULL;
    *count = 0;
    written = snprintf(path, sizeof path, "%s/tile_ra%03d_dec%03d.csv",
                       catalog->all_sky_root, ra_index, dec_index);
    if (written < 0 || (size_t)written >= sizeof path)
        return -1;
    StarCatalog temporary = {0};
    if (catalog_load_csv(&temporary, path))
        return -1;
    *stars = temporary.stars;
    *count = temporary.count;
    return 0;
}

static int load_tile(StarCatalog *catalog, int ra_index, int dec_index)
{
    CatalogTile *tile = &catalog->tiles[tile_index(ra_index, dec_index)];
    Star *stars;
    size_t count;
    if (tile->state != 0)
        return tile->state == 1 ? 0 : -1;
    if (load_tile_file(catalog, ra_index, dec_index, &stars, &count)) {
        tile->state = -1;
        return -1;
    }
    tile->stars = stars;
    tile->count = count;
    tile->state = 1;
    catalog->count += tile->count;
    return 0;
}

int catalog_load_all_sky(StarCatalog *catalog, const char *directory)
{
    size_t tile_count = (size_t)CATALOG_ALL_SKY_RA_TILES *
                        CATALOG_ALL_SKY_DEC_TILES;
    if (catalog == NULL || directory == NULL || directory[0] == '\0')
        return -1;
    *catalog = (StarCatalog){0};
    catalog->all_sky_root = malloc(strlen(directory) + 1);
    catalog->tiles = calloc(tile_count, sizeof *catalog->tiles);
    if (catalog->all_sky_root == NULL || catalog->tiles == NULL) {
        catalog_destroy(catalog);
        return -1;
    }
    strcpy(catalog->all_sky_root, directory);
    catalog->kind = STAR_CATALOG_ALL_SKY;
    return 0;
}

static void lon_lat_from_direction(const double direction[3], double *longitude,
                                   double *latitude)
{
    *longitude = atan2(direction[1], direction[0]) * 180.0 / PI;
    if (*longitude < 0.0)
        *longitude += 360.0;
    *latitude = asin(fmax(-1.0, fmin(1.0, direction[2]))) * 180.0 / PI;
}

static double dot(const double a[3], const double b[3])
{
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
}

static void cross(const double a[3], const double b[3], double out[3])
{
    out[0] = a[1] * b[2] - a[2] * b[1];
    out[1] = a[2] * b[0] - a[0] * b[2];
    out[2] = a[0] * b[1] - a[1] * b[0];
}

static void direction_from_lon_lat(double longitude, double latitude,
                                   double direction[3])
{
    const double lon = longitude * PI / 180.0;
    const double lat = latitude * PI / 180.0;
    const double cos_lat = cos(lat);
    direction[0] = cos_lat * cos(lon);
    direction[1] = cos_lat * sin(lon);
    direction[2] = sin(lat);
}

/* The minimum of a plane dot-product over a longitude/latitude rectangle is
 * attained on a boundary or at its antipodal stationary point.  This gives a
 * conservative, analytic whole-tile containment test; it is not a corner-only
 * approximation. */
static double tile_plane_minimum(const double normal[3], double lon_lo,
                                 double lon_hi, double lat_lo, double lat_hi)
{
    double values[32];
    size_t count = 0;
    const double phase = atan2(normal[2], normal[0]) * 180.0 / PI;
    const double latitude_phase = atan2(normal[1], hypot(normal[0], normal[2])) * 180.0 / PI;
    const double candidate_lon[] = {lon_lo, lon_hi, phase + 180.0, phase - 180.0};
    const double candidate_lat[] = {lat_lo, lat_hi, latitude_phase + 180.0,
                                    latitude_phase - 180.0};
    for (size_t i = 0; i < sizeof candidate_lon / sizeof *candidate_lon; ++i)
        for (size_t j = 0; j < 2; ++j) {
            double lon = candidate_lon[i];
            while (lon < lon_lo) lon += 360.0;
            while (lon > lon_hi) lon -= 360.0;
            if (lon >= lon_lo && lon <= lon_hi) {
                double point[3];
                direction_from_lon_lat(lon, j == 0 ? lat_lo : lat_hi, point);
                values[count++] = dot(normal, point);
            }
        }
    for (size_t i = 0; i < sizeof candidate_lon / sizeof *candidate_lon; ++i)
        for (size_t j = 0; j < sizeof candidate_lat / sizeof *candidate_lat; ++j) {
            double lon = candidate_lon[i];
            while (lon < lon_lo) lon += 360.0;
            while (lon > lon_hi) lon -= 360.0;
            if (lon >= lon_lo && lon <= lon_hi &&
                candidate_lat[j] >= lat_lo && candidate_lat[j] <= lat_hi) {
                double point[3];
                direction_from_lon_lat(lon, candidate_lat[j], point);
                values[count++] = dot(normal, point);
            }
        }
    for (size_t i = 0; i < 2; ++i)
        for (size_t j = 0; j < sizeof candidate_lat / sizeof *candidate_lat; ++j)
            if (candidate_lat[j] >= lat_lo && candidate_lat[j] <= lat_hi) {
                double point[3];
                direction_from_lon_lat(i == 0 ? lon_lo : lon_hi,
                                       candidate_lat[j], point);
                values[count++] = dot(normal, point);
            }
    double minimum = values[0];
    for (size_t i = 1; i < count; ++i)
        if (values[i] < minimum) minimum = values[i];
    return minimum;
}

static int tile_is_fully_contained(const double direction[3][3], int ra_index,
                                   int dec_index)
{
    const double lon_lo = ra_index;
    const double lon_hi = ra_index + 1.0;
    const double lat_lo = dec_index - 90.0;
    const double lat_hi = lat_lo + 1.0;
    for (int edge = 0; edge < 3; ++edge) {
        const double *left = direction[edge];
        const double *right = direction[(edge + 1) % 3];
        const double *opposite = direction[(edge + 2) % 3];
        double normal[3];
        cross(left, right, normal);
        if (dot(normal, opposite) < 0.0)
            for (int axis = 0; axis < 3; ++axis) normal[axis] = -normal[axis];
        if (tile_plane_minimum(normal, lon_lo, lon_hi, lat_lo, lat_hi) < -1e-14)
            return 0;
    }
    return 1;
}

typedef int (*CatalogTileIndexVisitor)(int ra_index, int dec_index,
                                       void *context);

static int visit_source_triangle_tile_indices(
    const double direction[3][3], CatalogTileIndexVisitor visitor,
    void *context)
{
    if (direction == NULL || visitor == NULL)
        return -1;
    double longitude[3], latitude[3], unwrapped[3];
    for (int i = 0; i < 3; ++i) {
        lon_lat_from_direction(direction[i], &longitude[i], &latitude[i]);
        unwrapped[i] = longitude[i];
        while (unwrapped[i] - longitude[0] > 180.0) unwrapped[i] -= 360.0;
        while (unwrapped[i] - longitude[0] < -180.0) unwrapped[i] += 360.0;
    }
    double lon_min = unwrapped[0], lon_max = unwrapped[0];
    double lat_min = latitude[0], lat_max = latitude[0];
    for (int i = 1; i < 3; ++i) {
        if (unwrapped[i] < lon_min) lon_min = unwrapped[i];
        if (unwrapped[i] > lon_max) lon_max = unwrapped[i];
        if (latitude[i] < lat_min) lat_min = latitude[i];
        if (latitude[i] > lat_max) lat_max = latitude[i];
    }
    /* A triangle containing a pole covers every RA there. */
    const double north[3] = {0.0, 1.0, 0.0};
    const double south[3] = {0.0, -1.0, 0.0};
    int all_ra = 0;
    for (int pole = 0; pole < 2; ++pole) {
        const double *point = pole == 0 ? north : south;
        int inside = 1;
        for (int edge = 0; edge < 3; ++edge) {
            double normal[3];
            cross(direction[edge], direction[(edge + 1) % 3], normal);
            if (dot(normal, point) * dot(normal, direction[(edge + 2) % 3]) < -1e-14) inside = 0;
        }
        if (inside) all_ra = 1;
    }
    /* Source edges are great-circle arcs, so their RA/Dec extrema need not be
     * vertices.  This small guard band covers that curvature without pulling
     * in an otherwise unrelated one-degree tile ring. */
    lon_min -= 0.01;
    lon_max += 0.01;
    lat_min -= 0.01;
    lat_max += 0.01;
    if (lon_max - lon_min >= 360.0)
        all_ra = 1;
    const int dec_first = fmax(0, (int)floor(lat_min + 90.0));
    const int dec_last = fmin(CATALOG_ALL_SKY_DEC_TILES - 1,
                              (int)floor(lat_max + 90.0));
    const int ra_first = (int)floor(lon_min);
    const int ra_last = (int)floor(lon_max);
    for (int dec = dec_first; dec <= dec_last; ++dec)
        for (int raw_ra = all_ra ? 0 : ra_first;
             raw_ra <= (all_ra ? 359 : ra_last); ++raw_ra) {
            const int ra = (raw_ra % 360 + 360) % 360;
            if (visitor(ra, dec, context))
                return -1;
        }
    return 0;
}

typedef struct {
    StarCatalog *catalog;
    const double (*direction)[3];
    int load_missing;
    CatalogTileVisitor visitor;
    void *context;
} CatalogVisitContext;

static int visit_catalog_tile(int ra_index, int dec_index, void *opaque)
{
    CatalogVisitContext *context = opaque;
    CatalogTile *tile =
        &context->catalog->tiles[tile_index(ra_index, dec_index)];
    if (context->load_missing && load_tile(context->catalog, ra_index, dec_index))
        return 0; /* Downloader has not finished this tile yet. */
    if (tile->state != 1)
        return 0;
    return context->visitor(
        tile->stars, tile->count,
        tile_is_fully_contained(context->direction, ra_index, dec_index),
        context->context);
}

typedef struct {
    unsigned char *requested;
} CatalogMarkContext;

static int mark_catalog_tile(int ra_index, int dec_index, void *opaque)
{
    CatalogMarkContext *context = opaque;
    context->requested[tile_index(ra_index, dec_index)] = 1;
    return 0;
}

int catalog_mark_source_triangle_tiles(
    const double direction[3][3],
    unsigned char requested[CATALOG_ALL_SKY_TILE_COUNT])
{
    CatalogMarkContext context = {.requested = requested};
    if (requested == NULL)
        return -1;
    return visit_source_triangle_tile_indices(direction, mark_catalog_tile,
                                              &context);
}

void catalog_tile_set_clear(CatalogTileSet *set)
{
    if (set != NULL)
        memset(set, 0, sizeof *set);
}

size_t catalog_tile_set_count(const CatalogTileSet *set)
{
    size_t count = 0;
    if (set == NULL)
        return 0;
    for (size_t tile_id = 0; tile_id < CATALOG_ALL_SKY_TILE_COUNT; ++tile_id)
        if (set->requested[tile_id])
            ++count;
    return count;
}

typedef struct {
    size_t tile_id;
    Star *stars;
    size_t count;
    int loaded;
} CatalogPendingTile;

/* Reads and commits up to `batch_tiles` unseen tiles from `ids` at a time.
 * Every batch allocates only its own bounded pending array, so an all-sky
 * union cannot stage an unbounded temporary copy of the star data. */
static int prefetch_tile_ids(StarCatalog *catalog, const size_t *ids,
                             size_t id_count, int worker_count,
                             size_t batch_tiles, CatalogPrefetchStats *stats)
{
    if (batch_tiles == 0)
        batch_tiles = CATALOG_PREFETCH_DEFAULT_BATCH_TILES;
    if (batch_tiles > id_count)
        batch_tiles = id_count;
    for (size_t base = 0; base < id_count; base += batch_tiles) {
        const size_t chunk = id_count - base < batch_tiles
                                 ? id_count - base
                                 : batch_tiles;
        CatalogPendingTile *pending = calloc(chunk, sizeof *pending);
        if (pending == NULL)
            return -1;
        size_t pending_count = 0;
        for (size_t i = 0; i < chunk; ++i) {
            const size_t tile_id = ids[base + i];
            if (catalog->tiles[tile_id].state != 0)
                continue;
            pending[pending_count++].tile_id = tile_id;
        }
        if (pending_count > 0) {
            int batch_workers = worker_count;
            if (batch_workers > (int)pending_count)
                batch_workers = (int)pending_count;
            const double load_start = omp_get_wtime();
#pragma omp parallel for num_threads(batch_workers) schedule(static)
            for (size_t i = 0; i < pending_count; ++i) {
                const int ra_index =
                    (int)(pending[i].tile_id % CATALOG_ALL_SKY_RA_TILES);
                const int dec_index =
                    (int)(pending[i].tile_id / CATALOG_ALL_SKY_RA_TILES);
                pending[i].loaded = !load_tile_file(
                    catalog, ra_index, dec_index, &pending[i].stars,
                    &pending[i].count);
            }
            if (stats != NULL)
                stats->load_seconds += omp_get_wtime() - load_start;
            /* Only this serial commit mutates the shared catalog cache. */
            for (size_t i = 0; i < pending_count; ++i) {
                CatalogTile *tile = &catalog->tiles[pending[i].tile_id];
                if (pending[i].loaded) {
                    tile->stars = pending[i].stars;
                    tile->count = pending[i].count;
                    tile->state = 1;
                    catalog->count += tile->count;
                    if (stats != NULL) {
                        ++stats->newly_loaded_tiles;
                        stats->newly_loaded_stars += tile->count;
                    }
                } else {
                    tile->state = -1;
                    if (stats != NULL)
                        ++stats->unavailable_tiles;
                }
            }
        }
        free(pending);
    }
    return 0;
}

/* Collects the requested-but-unseen tile ids, then reads them in batches. */
static int prefetch_bitmap(StarCatalog *catalog,
                           const unsigned char requested[CATALOG_ALL_SKY_TILE_COUNT],
                           int worker_count, size_t batch_tiles,
                           CatalogPrefetchStats *stats)
{
    size_t unseen = 0;
    if (stats != NULL)
        *stats = (CatalogPrefetchStats){0};
    if (catalog == NULL || requested == NULL || worker_count <= 0)
        return -1;
    if (catalog->kind != STAR_CATALOG_ALL_SKY)
        return 0;
    for (size_t tile_id = 0; tile_id < CATALOG_ALL_SKY_TILE_COUNT; ++tile_id) {
        if (!requested[tile_id])
            continue;
        if (stats != NULL)
            ++stats->requested_tiles;
        if (catalog->tiles[tile_id].state == 0)
            ++unseen;
    }
    if (unseen == 0)
        return 0;
    size_t *ids = malloc(unseen * sizeof *ids);
    if (ids == NULL)
        return -1;
    size_t index = 0;
    for (size_t tile_id = 0; tile_id < CATALOG_ALL_SKY_TILE_COUNT; ++tile_id)
        if (requested[tile_id] && catalog->tiles[tile_id].state == 0)
            ids[index++] = tile_id;
    const int result =
        prefetch_tile_ids(catalog, ids, unseen, worker_count, batch_tiles, stats);
    free(ids);
    return result;
}

int catalog_prefetch_marked_tiles(
    StarCatalog *catalog,
    const unsigned char requested[CATALOG_ALL_SKY_TILE_COUNT],
    int worker_count, CatalogPrefetchStats *stats)
{
    return prefetch_bitmap(catalog, requested, worker_count,
                           CATALOG_ALL_SKY_TILE_COUNT, stats);
}

int catalog_prefetch_tile_set(StarCatalog *catalog, const CatalogTileSet *set,
                              int worker_count, size_t batch_tiles,
                              CatalogPrefetchStats *stats)
{
    if (set == NULL)
        return -1;
    return prefetch_bitmap(catalog, set->requested, worker_count, batch_tiles,
                           stats);
}

int catalog_visit_source_triangle(StarCatalog *catalog,
                                  const double direction[3][3],
                                  int load_missing, CatalogTileVisitor visitor,
                                  void *context)
{
    if (catalog == NULL || direction == NULL || visitor == NULL)
        return -1;
    if (catalog->kind == STAR_CATALOG_MEMORY)
        return visitor(catalog->stars, catalog->count, 0, context);
    CatalogVisitContext visit_context = {
        .catalog = catalog, .direction = direction, .load_missing = load_missing,
        .visitor = visitor, .context = context};
    return visit_source_triangle_tile_indices(direction, visit_catalog_tile,
                                              &visit_context);
}

void catalog_destroy(StarCatalog *catalog)
{
    if (catalog->tiles != NULL)
        for (size_t i = 0;
             i < (size_t)CATALOG_ALL_SKY_RA_TILES * CATALOG_ALL_SKY_DEC_TILES;
             ++i)
            free(catalog->tiles[i].stars);
    free(catalog->stars);
    free(catalog->tiles);
    free(catalog->all_sky_root);
    catalog->stars = NULL;
    catalog->count = 0;
    catalog->tiles = NULL;
    catalog->all_sky_root = NULL;
    catalog->kind = STAR_CATALOG_MEMORY;
}