/**
 * @file cif.c
 * @brief mmCIF parsing pipeline.
 *
 * Parses mmCIF files to extract molecular structure data including
 * coordinates, atom types, residue sequences, and chain organization.
 *
 * Pipeline:
 *   1. Block Validation  - Verify required mmCIF blocks exist
 *   2. Metadata Parsing  - Extract chain/residue counts and names
 *   3. Atom Parsing      - Single-pass coordinate/type extraction (parallelized)
 *   4. Atom Reordering   - Separate polymer/non-polymer atoms
 */

#include "cif.h"

#ifdef _OPENMP
#include <omp.h>
#endif

/* Hash tables for type lookups */
#include "hash/atom.c"
#include "hash/residue.c"
#include "hash/element.c"


/* ============================================================================
 * CONSTANTS
 * mmCIF attribute names used throughout parsing.
 * ============================================================================ */

static const size_t COORDS = 3;

/* Coordinate attributes */
static const char *ATTR_X = "Cartn_x";
static const char *ATTR_Y = "Cartn_y";
static const char *ATTR_Z = "Cartn_z";

/* Structure attributes */
static const char *ATTR_MODEL         = "pdbx_PDB_model_num";
static const char *ATTR_CHAIN_ID      = "id";
static const char *ATTR_RES_PER_CHAIN = "asym_id";
static const char *ATTR_STRAND_ID     = "pdb_strand_id";
static const char *ATTR_RESIDUE_NAME  = "mon_id";
static const char *ATTR_ELEMENT       = "type_symbol";
static const char *ATTR_ATOM_NAME     = "label_atom_id";
static const char *ATTR_SEQ_ID        = "label_seq_id";
static const char *ATTR_LABEL_ASYM    = "label_asym_id";
static const char *ATTR_COMP_ID       = "label_comp_id";


/* ============================================================================
 * TYPES
 * Internal structures for parsing state.
 * ============================================================================ */

/**
 * Pre-computed attribute indices for single-pass atom parsing.
 *
 * Storing indices avoids repeated string lookups in the hot loop.
 * Extend this struct to add new per-atom features (b-factor, occupancy, etc).
 */
typedef struct {
    int x, y, z;              /* Coordinate column indices */
    int element;              /* Element symbol (C, N, O, etc.) */
    int comp_id;              /* Residue/component name */
    int atom_name;            /* Atom name within residue */
    int seq_id;               /* Sequence ID for polymer detection */
    int label_asym;           /* Chain label */
} AtomIndices;


/* ============================================================================
 * FILE HEADER
 * Extract PDB identifier from mmCIF header.
 * ============================================================================ */

/**
 * Parse the PDB identifier from "data_XXXX" header line.
 */
char *_get_id(char *buffer, CifErrorContext *ctx) {
    const char *prefix = "data_";

    if (_neq(buffer, prefix)) {
        CIF_SET_ERROR(ctx, CIF_ERR_PARSE,
            "Invalid mmCIF file: missing 'data_' prefix");
        return NULL;
    }

    buffer += 5;
    char *start = buffer;
    while (*buffer != '\n' && *buffer != '\0') buffer++;

    return _strdup_n(start, (size_t)(buffer - start), ctx);
}


/* ============================================================================
 * ATTRIBUTE UTILITIES
 * Generic helpers for extracting values from mmCIF blocks.
 * ============================================================================ */

/**
 * Extract unique string values from an attribute (first occurrence only).
 *
 * Used for chain names, strand IDs, etc. where values repeat across rows
 * but we only want distinct values in order of appearance.
 */
static char **_get_unique(mmBlock *block, const char *attr, int *size,
                          CifErrorContext *ctx) {
    int index = _get_attr_index(block, attr);
    if (index == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR,
            "Missing attribute '%s' in block '%s'", attr, block->category);
        return NULL;
    }

    size_t alloc_size = (size_t)(*size > 0 ? *size : block->size);
    char **str = calloc(alloc_size, sizeof(char *));
    if (str == NULL) {
        CIF_SET_ERROR(ctx, CIF_ERR_ALLOC,
            "Failed to allocate unique array of size %zu", alloc_size);
        return NULL;
    }

    char *prev = NULL;
    int ix = 0;

    for (int line = 0; line < block->size; line++) {
        char *token = _get_attr_by_line(block, line, index, ctx);
        if (token == NULL) {
            for (int i = 0; i <= ix; i++) free(str[i]);
            free(str);
            return NULL;
        }

        if (prev == NULL) {
            prev = token;
            str[ix] = token;
        } else if (_neq(prev, token)) {
            prev = token;
            if ((size_t)++ix >= alloc_size) {
                free(token);
                break;
            }
            str[ix] = token;
        } else {
            free(token);
        }
    }

    if (*size <= 0) {
        *size = ix + 1;
        char **resized = realloc(str, (size_t)(*size) * sizeof(char *));
        return resized ? resized : str;
    }
    return str;
}

/**
 * Count unique consecutive values in an attribute.
 */
static int _count_unique(mmBlock *block, const char *attr, CifErrorContext *ctx) {
    int index = _get_attr_index(block, attr);
    if (index == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR,
            "Missing attribute '%s' in block '%s'", attr, block->category);
        return -1;
    }

    int count = 0;
    char *prev = NULL;

    for (int line = 0; line < block->size; line++) {
        char *token = _get_attr_by_line(block, line, index, ctx);
        if (token == NULL) {
            free(prev);
            return -1;
        }

        if (prev == NULL || _neq(prev, token)) {
            free(prev);
            prev = token;
            count++;
        } else {
            free(token);
        }
    }

    free(prev);
    return count;
}

/**
 * Parse residue types via hash table lookup.
 *
 * Used for sequence parsing where residue names map to type indices.
 */
static int *_parse_via_lookup(mmBlock *block, HashTable func, const char *attr,
                              CifErrorContext *ctx) {
    int index = _get_attr_index(block, attr);
    if (index == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR,
            "Missing attribute '%s' in block '%s'", attr, block->category);
        return NULL;
    }

    int *array = calloc((size_t)block->size, sizeof(int));
    if (array == NULL) {
        CIF_SET_ERROR(ctx, CIF_ERR_ALLOC,
            "Failed to allocate lookup array of size %d", block->size);
        return NULL;
    }

    for (int line = 0; line < block->size; line++) {
        char *token = _get_attr_by_line(block, line, index, ctx);
        if (token == NULL) {
            free(array);
            return NULL;
        }
        array[line] = _lookup(func, token);
        free(token);
    }

    return array;
}


/* ============================================================================
 * SIZE COUNTING
 * Count items per group (residues per chain, atoms per residue, etc).
 * ============================================================================ */

/**
 * Count items grouped by attribute value changes.
 *
 * Returns array where sizes[i] = count of rows with i-th unique value.
 * Used for residues-per-chain and atoms-per-chain counting.
 */
static int *_count_sizes_by_group(mmBlock *block, const char *attr, int *size,
                                  CifErrorContext *ctx) {
    int index = _get_attr_index(block, attr);
    if (index == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR,
            "Missing attribute '%s' in block '%s'", attr, block->category);
        return NULL;
    }

    size_t alloc_size = (size_t)(*size > 0 ? *size : block->size);
    int *sizes = calloc(alloc_size, sizeof(int));
    if (sizes == NULL) {
        CIF_SET_ERROR(ctx, CIF_ERR_ALLOC,
            "Failed to allocate sizes array of size %zu", alloc_size);
        return NULL;
    }

    char *prev = NULL;
    int ix = 0;

    for (int line = 0; line < block->size; line++) {
        char *token = _get_attr_by_line(block, line, index, ctx);
        if (token == NULL) {
            free(prev);
            free(sizes);
            return NULL;
        }

        if (prev == NULL) {
            prev = token;
        } else if (_neq(prev, token)) {
            free(prev);
            prev = token;
            ix++;
        } else {
            free(token);
        }

        if ((size_t)ix < alloc_size) sizes[ix]++;
    }

    free(prev);

    if (*size <= 0) {
        *size = ix + 1;
        int *resized = realloc(sizes, (size_t)(*size) * sizeof(int));
        return resized ? resized : sizes;
    }
    return sizes;
}

/**
 * Count atoms per residue with chain-aware indexing.
 *
 * Handles non-polymer atoms (seq_id < 0) by marking them in is_nonpoly mask.
 * Returns NULL-terminated sizes array indexed by global residue index.
 */
static int *_count_atoms_per_residue(mmBlock *block, int residue_count,
                                     int *nonpoly_count, int *is_nonpoly,
                                     int *res_per_chain, CifErrorContext *ctx) {
    int seq_index = _get_attr_index(block, ATTR_SEQ_ID);
    if (seq_index == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR, "Missing attribute '%s'", ATTR_SEQ_ID);
        return NULL;
    }

    int chain_index = _get_attr_index(block, ATTR_LABEL_ASYM);
    if (chain_index == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR, "Missing attribute '%s'", ATTR_LABEL_ASYM);
        return NULL;
    }

    int *sizes = calloc((size_t)residue_count, sizeof(int));
    if (sizes == NULL) {
        CIF_SET_ERROR(ctx, CIF_ERR_ALLOC,
            "Failed to allocate residue sizes array of size %d", residue_count);
        return NULL;
    }

    int chain_offset = 0;
    char *prev_chain = NULL;
    int *chain_len_ptr = res_per_chain;

    for (int line = 0; line < block->size; line++) {
        /* Track chain changes to compute residue offset */
        char *chain = _get_attr_by_line(block, line, chain_index, ctx);
        if (chain == NULL) {
            free(prev_chain);
            free(sizes);
            return NULL;
        }

        if (prev_chain == NULL) {
            prev_chain = chain;
        } else if (_neq(prev_chain, chain)) {
            free(prev_chain);
            prev_chain = chain;
            chain_offset += *chain_len_ptr++;
        } else {
            free(chain);
        }

        /* Parse sequence ID */
        char *seq_token = _get_attr_by_line(block, line, seq_index, ctx);
        if (seq_token == NULL) {
            free(prev_chain);
            free(sizes);
            return NULL;
        }

        int seq_id = _str_to_int(seq_token) - 1;
        free(seq_token);

        /* Non-polymer atoms have invalid seq_id */
        if (seq_id < 0) {
            (*nonpoly_count)++;
            is_nonpoly[line] = 1;
            continue;
        }

        is_nonpoly[line] = 0;
        int residue_idx = chain_offset + seq_id;
        if (residue_idx >= 0 && residue_idx < residue_count) {
            sizes[residue_idx]++;
        }
    }

    free(prev_chain);
    return sizes;
}


/* ============================================================================
 * BATCH ATOM PARSING
 * Single-pass parallel parsing of all per-atom data.
 * Maximizes cache efficiency by processing each line once.
 * ============================================================================ */

/**
 * Initialize attribute indices for batch parsing.
 *
 * Pre-computes column indices to avoid repeated string lookups in hot loop.
 */
static CifError _init_atom_indices(mmBlock *block, AtomIndices *idx,
                                   CifErrorContext *ctx) {
    idx->x = _get_attr_index(block, ATTR_X);
    idx->y = _get_attr_index(block, ATTR_Y);
    idx->z = _get_attr_index(block, ATTR_Z);
    idx->element = _get_attr_index(block, ATTR_ELEMENT);
    idx->comp_id = _get_attr_index(block, ATTR_COMP_ID);
    idx->atom_name = _get_attr_index(block, ATTR_ATOM_NAME);
    idx->seq_id = _get_attr_index(block, ATTR_SEQ_ID);
    idx->label_asym = _get_attr_index(block, ATTR_LABEL_ASYM);

    /* Validate required indices */
    if (idx->x == BAD_IX || idx->y == BAD_IX || idx->z == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR, "Missing coordinate attributes");
        return CIF_ERR_ATTR;
    }
    if (idx->element == BAD_IX || idx->comp_id == BAD_IX || idx->atom_name == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR, "Missing atom type attributes");
        return CIF_ERR_ATTR;
    }
    if (idx->seq_id == BAD_IX || idx->label_asym == BAD_IX) {
        CIF_SET_ERROR(ctx, CIF_ERR_ATTR, "Missing chain/residue attributes");
        return CIF_ERR_ATTR;
    }

    return CIF_OK;
}

/**
 * Parse all per-atom data in a single parallel pass.
 *
 * For each atom line, parses:
 *   - Coordinates (x, y, z)
 *   - Element type via hash lookup
 *   - Atom type via combined hash lookup (residue_atom)
 *
 * Uses pre-computed line pointers for O(1) row access.
 * Unknown types return -1 (valid for non-standard residues/ligands).
 */
static CifError _parse_atoms_batch(mmBlock *block, AtomIndices *idx,
                                   float *coords, int *elements, int *types) {
    #pragma omp parallel
    {
        char combine_buf[MAX_INLINE_BUFFER];

        #pragma omp for schedule(static)
        for (int line = 0; line < block->size; line++) {
            /* Parse coordinates */
            coords[COORDS * line + 0] = _parse_float_inline(block, line, idx->x);
            coords[COORDS * line + 1] = _parse_float_inline(block, line, idx->y);
            coords[COORDS * line + 2] = _parse_float_inline(block, line, idx->z);

            /* Parse types (-1 for unknown) */
            elements[line] = _lookup_inline(block, line, idx->element, _lookup_element);
            types[line] = _lookup_double_inline(block, line, idx->comp_id,
                                                idx->atom_name, _lookup_atom, combine_buf);
        }
    }

    return CIF_OK;
}


/* ============================================================================
 * ATOM REORDERING
 * Separate polymer atoms from non-polymer (ligands, water, ions).
 * Enables simple slicing: polymer = atoms[0:polymer_count]
 * ============================================================================ */

/**
 * Reorder arrays so polymer atoms come first.
 *
 * After reordering:
 *   - atoms[0, polymer_count) = polymer atoms
 *   - atoms[polymer_count, total) = non-polymer atoms
 */
static CifError _reorder_atoms(mmCIF *cif, int *is_nonpoly, CifErrorContext *ctx) {
    int atoms = cif->atoms;
    int polymer_count = atoms - cif->nonpoly;

    if (cif->nonpoly == 0) {
        cif->polymer = polymer_count;
        return CIF_OK;
    }

    /* Allocate reorder buffers */
    float *new_coords = calloc(COORDS * (size_t)atoms, sizeof(float));
    int *new_types = calloc((size_t)atoms, sizeof(int));
    int *new_elements = calloc((size_t)atoms, sizeof(int));

    if (!new_coords || !new_types || !new_elements) {
        free(new_coords);
        free(new_types);
        free(new_elements);
        CIF_SET_ERROR(ctx, CIF_ERR_ALLOC, "Failed to allocate reorder buffers");
        return CIF_ERR_ALLOC;
    }

    /* Two-pointer copy: polymer to front, non-polymer to back */
    int poly_ix = 0;
    int nonpoly_ix = polymer_count;

    for (int i = 0; i < atoms; i++) {
        int dest = is_nonpoly[i] ? nonpoly_ix++ : poly_ix++;
        memcpy(&new_coords[COORDS * dest], &cif->coordinates[COORDS * i],
               COORDS * sizeof(float));
        new_types[dest] = cif->types[i];
        new_elements[dest] = cif->elements[i];
    }

    /* Swap arrays */
    free(cif->coordinates);
    free(cif->types);
    free(cif->elements);

    cif->coordinates = new_coords;
    cif->types = new_types;
    cif->elements = new_elements;
    cif->polymer = polymer_count;

    return CIF_OK;
}


/* ============================================================================
 * PUBLIC INTERFACE
 * Main entry point for populating mmCIF structure from parsed blocks.
 * ============================================================================ */

/**
 * Populate mmCIF structure from parsed blocks.
 *
 * Pipeline:
 *   1. Validate required blocks exist
 *   2. Count models, chains, residues, atoms
 *   3. Parse chain/residue metadata
 *   4. Batch parse atom data (parallelized)
 *   5. Reorder atoms (polymer first)
 */
CifError _fill_cif(mmCIF *cif, mmBlockList *blocks, CifErrorContext *ctx) {

    /* ── Block Validation ─────────────────────────────────────────────────── */

    if (blocks->atom.category == NULL) {
        CIF_SET_ERROR(ctx, CIF_ERR_BLOCK, "Missing required _atom_site block");
        return CIF_ERR_BLOCK;
    }
    if (blocks->poly.category == NULL) {
        CIF_SET_ERROR(ctx, CIF_ERR_BLOCK, "Missing required _pdbx_poly_seq_scheme block");
        return CIF_ERR_BLOCK;
    }
    if (blocks->chain.category == NULL) {
        CIF_SET_ERROR(ctx, CIF_ERR_BLOCK, "Missing required _struct_asym block");
        return CIF_ERR_BLOCK;
    }

    /* ── Count Structure Elements ─────────────────────────────────────────── */

    int model_count = _count_unique(&blocks->atom, ATTR_MODEL, ctx);
    if (model_count < 0) return ctx->code;
    cif->models = model_count;

    cif->chains = blocks->chain.size;
    cif->residues = blocks->poly.size;

    /* Adjust for multi-model structures (use first model only) */
    blocks->atom.size /= cif->models;
    cif->atoms = blocks->atom.size;

    /* ── Parse Metadata ───────────────────────────────────────────────────── */

    cif->res_per_chain = _count_sizes_by_group(&blocks->poly, ATTR_RES_PER_CHAIN,
                                               &cif->chains, ctx);
    if (cif->res_per_chain == NULL) return ctx->code;

    cif->names = _get_unique(&blocks->chain, ATTR_CHAIN_ID, &cif->chains, ctx);
    if (cif->names == NULL) return ctx->code;

    cif->sequence = _parse_via_lookup(&blocks->poly, _lookup_residue,
                                      ATTR_RESIDUE_NAME, ctx);
    if (cif->sequence == NULL) return ctx->code;

    cif->strands = _get_unique(&blocks->poly, ATTR_STRAND_ID, &cif->chains, ctx);
    if (cif->strands == NULL) return ctx->code;

    /* ── Batch Atom Parsing ───────────────────────────────────────────────── */

    CifError err = _precompute_lines(&blocks->atom, ctx);
    if (err != CIF_OK) return err;

    AtomIndices idx;
    err = _init_atom_indices(&blocks->atom, &idx, ctx);
    if (err != CIF_OK) {
        _free_lines(&blocks->atom);
        return err;
    }

    cif->coordinates = calloc(COORDS * (size_t)cif->atoms, sizeof(float));
    cif->elements = calloc((size_t)cif->atoms, sizeof(int));
    cif->types = calloc((size_t)cif->atoms, sizeof(int));

    if (!cif->coordinates || !cif->elements || !cif->types) {
        free(cif->coordinates);
        free(cif->elements);
        free(cif->types);
        _free_lines(&blocks->atom);
        CIF_SET_ERROR(ctx, CIF_ERR_ALLOC, "Failed to allocate atom arrays");
        return CIF_ERR_ALLOC;
    }

    err = _parse_atoms_batch(&blocks->atom, &idx,
                             cif->coordinates, cif->elements, cif->types);
    _free_lines(&blocks->atom);
    if (err != CIF_OK) {
        free(cif->coordinates);
        free(cif->elements);
        free(cif->types);
        return err;
    }

    /* ── Atom Reordering ──────────────────────────────────────────────────── */

    int *is_nonpoly = calloc((size_t)cif->atoms, sizeof(int));
    if (is_nonpoly == NULL) {
        CIF_SET_ERROR(ctx, CIF_ERR_ALLOC, "Failed to allocate nonpoly mask");
        return CIF_ERR_ALLOC;
    }

    cif->nonpoly = 0;
    cif->atoms_per_res = _count_atoms_per_residue(&blocks->atom, cif->residues,
                                                  &cif->nonpoly, is_nonpoly,
                                                  cif->res_per_chain, ctx);
    if (cif->atoms_per_res == NULL) {
        free(is_nonpoly);
        return ctx->code;
    }

    err = _reorder_atoms(cif, is_nonpoly, ctx);
    free(is_nonpoly);
    if (err != CIF_OK) return err;

    cif->atoms_per_chain = _count_sizes_by_group(&blocks->atom, ATTR_LABEL_ASYM,
                                                 &cif->chains, ctx);
    if (cif->atoms_per_chain == NULL) return ctx->code;

    return CIF_OK;
}
