/**
 * @file c_example.c
 * @brief Example usage of TACO format C API
 * 
 * This example demonstrates how to write and read TACO trajectory files
 * using the C API.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "taco_format.h"

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

// Simple water molecule example (3 atoms: O, H, H)
#define NUM_ATOMS 3
#define NUM_FRAMES 100

int write_example() {
    printf("Writing TACO file example...\n");

    // Set up atom metadata
    float masses[NUM_ATOMS] = {15.999, 1.008, 1.008}; // O, H, H
    const char* atom_names[NUM_ATOMS] = {"O", "H1", "H2"};
    const char* atom_types[NUM_ATOMS] = {"O", "H", "H"};

    taco_atom_metadata_t atom_metadata = {
        .masses = masses,
        .names = atom_names,
        .types = atom_types,
        .num_atoms = NUM_ATOMS
    };

    // Set up simulation metadata
    const char* ensemble = "NVT";
    taco_simulation_metadata_t sim_metadata = {
        .ensemble = ensemble,
        .temperature = 300.0,
        .pressure = 1.0,
        .num_degrees_freedom = 9
    };

    // Set up compression settings
    taco_compression_settings_t compression = {
        .position_precision = 1,  // Half precision
        .velocity_precision = 1,  // Half precision
        .force_precision = 0,     // Lossless
        .compression_level = 3
    };

    // Create writer
    taco_writer_t* writer = taco_writer_create("example.taco", NUM_ATOMS, 0.001, 
                                             &sim_metadata, &atom_metadata, compression);
    if (!writer) {
        fprintf(stderr, "Failed to create TACO writer\n");
        return -1;
    }

    // Write frames
    for (int frame = 0; frame < NUM_FRAMES; frame++) {
        taco_frame_t taco_frame = {0};
        taco_frame.frame_number = frame;
        taco_frame.time = frame * 0.001; // 1 fs time step
        taco_frame.num_atoms = NUM_ATOMS;
        
        // Generate simple water molecule positions
        float positions[NUM_ATOMS * 3];
        float velocities[NUM_ATOMS * 3];
        float forces[NUM_ATOMS * 3];
        
        // Simple animation: water molecule vibrating
        double t = frame * 0.001;
        
        // Oxygen at origin with small displacement
        positions[0] = 0.1 * sin(2 * M_PI * t);  // x
        positions[1] = 0.0;                      // y
        positions[2] = 0.0;                      // z
        
        // First hydrogen
        positions[3] = 0.1 + 0.05 * cos(2 * M_PI * t);  // x
        positions[4] = 0.08;                             // y
        positions[5] = 0.0;                              // z
        
        // Second hydrogen
        positions[6] = 0.1 + 0.05 * cos(2 * M_PI * t);  // x
        positions[7] = -0.08;                            // y
        positions[8] = 0.0;                              // z
        
        // Simple velocities (derivatives of positions)
        for (int i = 0; i < NUM_ATOMS * 3; i++) {
            velocities[i] = 0.1 * (rand() / (float)RAND_MAX - 0.5);
            forces[i] = 0.01 * (rand() / (float)RAND_MAX - 0.5);
        }
        
        // Box dimensions (cubic box)
        float box[3] = {2.0, 2.0, 2.0};
        
        taco_frame.positions = positions;
        taco_frame.velocities = velocities;
        taco_frame.forces = forces;
        taco_frame.box_dimensions = box;
        taco_frame.box_angles = NULL; // Default orthogonal
        taco_frame.potential_energy = -10.0 - 0.1 * sin(2 * M_PI * t);
        taco_frame.kinetic_energy = 0.5 + 0.1 * cos(2 * M_PI * t);
        taco_frame.temperature = 295.0 + 10.0 * sin(2 * M_PI * t);
        
        int result = taco_writer_write_frame(writer, &taco_frame);
        if (result != TACO_SUCCESS) {
            fprintf(stderr, "Failed to write frame %d: %d\n", frame, result);
            taco_writer_free(writer);
            return -1;
        }
    }

    // Finish writing
    int result = taco_writer_finish(writer);
    if (result != TACO_SUCCESS) {
        fprintf(stderr, "Failed to finish writing: %d\n", result);
        return -1;
    }

    printf("Successfully wrote %d frames to example.taco\n", NUM_FRAMES);
    return 0;
}

int read_example() {
    printf("Reading TACO file example...\n");

    // Open file for reading
    taco_reader_t* reader = taco_reader_open("example.taco");
    if (!reader) {
        fprintf(stderr, "Failed to open TACO file for reading\n");
        return -1;
    }

    uint64_t num_frames = taco_reader_num_frames(reader);
    uint32_t num_atoms = taco_reader_num_atoms(reader);

    printf("File contains %lu frames with %u atoms each\n", num_frames, num_atoms);

    // Read first and last frames
    for (uint64_t frame_idx = 0; frame_idx < num_frames; frame_idx += (num_frames - 1)) {
        taco_frame_t frame = {0};
        int result = taco_reader_read_frame(reader, frame_idx, &frame);
        
        if (result != TACO_SUCCESS) {
            fprintf(stderr, "Failed to read frame %lu: %d\n", frame_idx, result);
            continue;
        }

        printf("\nFrame %lu (time = %.3f ps):\n", frame.frame_number, frame.time);
        printf("  Potential energy: %.3f kJ/mol\n", frame.potential_energy);
        printf("  Kinetic energy: %.3f kJ/mol\n", frame.kinetic_energy);
        printf("  Temperature: %.1f K\n", frame.temperature);
        
        if (frame.box_dimensions) {
            printf("  Box: [%.3f, %.3f, %.3f] nm\n", 
                   frame.box_dimensions[0], frame.box_dimensions[1], frame.box_dimensions[2]);
        }

        // Print first atom position
        if (frame.positions) {
            printf("  First atom position: [%.4f, %.4f, %.4f] nm\n",
                   frame.positions[0], frame.positions[1], frame.positions[2]);
        }

        // Free frame arrays
        taco_frame_free_arrays(&frame);
    }

    taco_reader_free(reader);
    printf("Successfully read TACO file\n");
    return 0;
}

int main() {
    printf("TACO Format C API Example\n");
    printf("=========================\n\n");

    // Write example file
    if (write_example() != 0) {
        return 1;
    }

    printf("\n");

    // Read example file
    if (read_example() != 0) {
        return 1;
    }

    printf("\nExample completed successfully!\n");
    return 0;
}
