#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include "sphere_adapt_functions.h"

/* Function definition for magnetized_sphere_anomaly
   -calculate the magnetic anomaly due to a buried sphere. 
   Parameters  x = an observation point northing coordinate (m)
               y = an observation point easting coordinate (m)
               (x and y are relative to the sphere center, which is assumed to be (0,0))
   Sphere includes
               easting = easting coordinate of the sphere center (m)
               northing = northing coordinate of the sphere center (m)
               z = depth the center of the sphere (m)
               a = radius of the sphere
               minc = inclination of the vector of magnetization (deg)
               mdec = declination of the vector of magnetization (deg)
               mi = intensity of magnetization (amp/m)
               einc = inclination of the earths magnetic field (deg)
               edec = declination of the earths magnetic field (deg)
        return the total anomaly of the magnetic field
*/
double magnetized_sphere_anomaly(double x, double y, Sphere sphere) {
    double minc = sphere.minc * M_PI / 180.0; // Inclination down in rad
    double mdec = sphere.mdec * M_PI / 180.0; // Declination east in rad

    double ml = cos(minc) * cos(mdec);
    double mm = cos(minc) * sin(mdec);
    double mn = sin(minc);

    double dpm = (4.0 / 3.0) * M_PI * pow(sphere.a, 3) * sphere.mi;

    double mx = sphere.mi * ml;
    double my = sphere.mi * mm;
    double mz = sphere.mi * mn;

    double einc = sphere.einc * M_PI / 180.0; // Inclination down in rad
    double edec = sphere.edec * M_PI / 180.0; // Declination east in rad

    double el = cos(einc) * cos(edec);
    double em = cos(einc) * sin(edec);
    double en = sin(einc);

    double prop = 400 * M_PI;

    double z = - sphere.z;
    
    x = x - sphere.northing;
    y = y - sphere.easting;

    double r2 = x * x + y * y + z * z;
    double r = sqrt(r2);
    double r5 = pow(r, 5);

    double dotprod = x * mx + y * my + z * mz;

    double bx = prop * dpm * (3 * dotprod * x - r2 * mx) / r5;
    double by = prop * dpm * (3 * dotprod * y - r2 * my) / r5;
    double bz = prop * dpm * (3 * dotprod * z - r2 * mz) / r5;

    return el * bx + em * by + en * bz;
}


// Function prototype for calculating RMSE
double calculate_rmse(DataPoint *mag_anomaly, int size) {
    double sum = 0.0;
    for (int i = 0; i < size; i++) {
        double diff = mag_anomaly[i].mag_obs - mag_anomaly[i].mag_calc;
        sum += diff * diff;
    }
    return sqrt(sum / size);
}

double multistart_gradient_search(DataPoint *mag_anomaly, Sphere *sphere, int size, double alpha_init, double tol, int max_iter, int num_starts, double z_min, double z_max) {
    double best_z = z_min;
    double best_rmse = INFINITY;

    for (int i = 0; i < num_starts; i++) {
        // Generate an initial guess within the range [z_min, z_max]
        double initial_z = z_min + (z_max - z_min) * i / (num_starts - 1);
        sphere->z = initial_z;

        // Perform gradient search
        double optimized_z = adaptive_gradient_search(mag_anomaly, sphere, size, alpha_init, tol, max_iter);

        // Evaluate RMSE at the optimized depth
        sphere->z = optimized_z;
        for (int j = 0; j < size; j++) {
            mag_anomaly[j].mag_calc = magnetized_sphere_anomaly(mag_anomaly[j].northing, mag_anomaly[j].easting, *sphere);
        }
        double rmse = calculate_rmse(mag_anomaly, size);

        // Update the best result
        if (rmse < best_rmse) {
            best_rmse = rmse;
            best_z = optimized_z;
        }
    }

    return best_z; // Return the best depth found
}

double coarse_to_fine_search(DataPoint *mag_anomaly, Sphere *sphere, int size, double z_min, double z_max, int grid_points) {
    double best_z = z_min;
    double best_rmse = INFINITY;

    // Coarse search: grid search over depths
    double step = (z_max - z_min) / (float)(grid_points - 1);
    //printf("step = %lf\n", step);
    for (double z = z_min; z <= z_max; z += step) {
        sphere->z = z;
        //printf("sherez = %lf", sphere->z);
        // Evaluate RMSE
        for (int i = 0; i < size; i++) {
            mag_anomaly[i].mag_calc = magnetized_sphere_anomaly(mag_anomaly[i].northing, mag_anomaly[i].easting, *sphere);
        }
        double rmse = calculate_rmse(mag_anomaly, size);

        // Track the best depth
        if (rmse < best_rmse) {
            best_rmse = rmse;
            best_z = z;
        }
    }
        return best_z;
    }

double adaptive_gradient_search(DataPoint *mag_anomaly, Sphere *sphere, int size, double alpha_init, double tol, int max_iter) {
    double z = sphere->z;  // Initial depth guess
    double prev_z;
    double grad;
    double delta = 1.0; // Small step for numerical gradient calculation (meters depth)
    double alpha = alpha_init; // Initial learning rate
    double gamma = 0.9; // Decay factor for adaptive learning rate
    int iter = 0;

    do {
        prev_z = z;

        // Compute RMSE at z
        sphere->z = z;
        for (int i = 0; i < size; i++) {
            mag_anomaly[i].mag_calc = magnetized_sphere_anomaly(mag_anomaly[i].northing, mag_anomaly[i].easting, *sphere);
        }
        double rmse_z = calculate_rmse(mag_anomaly, size);

        // Compute RMSE at (z + delta) for numerical gradient
        sphere->z = z + delta;
        for (int i = 0; i < size; i++) {
            mag_anomaly[i].mag_calc = magnetized_sphere_anomaly(mag_anomaly[i].northing, mag_anomaly[i].easting, *sphere);
        }
        double rmse_z_plus = calculate_rmse(mag_anomaly, size);

        // Compute gradient (approximate derivative)
        grad = (rmse_z_plus - rmse_z) / delta;

        // Adaptive learning rate adjustment
        if (fabs(grad) > tol) {
            alpha *= gamma; // Reduce step size if gradient is large
        } else {
            alpha /= gamma; // Increase step size if gradient is small
        }

        // Prevent alpha from becoming too small or too large
        if (alpha > alpha_init) alpha = alpha_init; 
        if (alpha < 1e-6) alpha = 1e-6;

        // Update z using gradient descent
        z -= alpha * grad;

        iter++;

    } while (fabs(z - prev_z) > tol && fabs(grad) > tol && iter < max_iter);

    sphere->z = z; // Update the sphere's depth with optimized value
    return z; // Return the optimized depth
}

/* Function to read the data file
The data file must include four columns
northing (m))
easting (m))
observed magnetic value (nT))
calculated magnetic value (nT))

separated by blank space
The initial value of the calculated magnetic value might be 0.0
*/
int read_data_file(const char *filename, DataPoint **mag_anomaly, int *row_count) {
    FILE *file = fopen(filename, "r");
    if (file == NULL) {
        perror("Error opening data file");
        return -1;
    }

    char line[256];
    int capacity = 10; // Initial allocation size
    *row_count = 0;

    // Allocate initial memory for the mag_anomaly data struct
    *mag_anomaly = (DataPoint *)malloc(capacity * sizeof(DataPoint));
    

    if (*mag_anomaly == NULL) {
        perror("Error allocating memory");
        fclose(file);
        return -1;
    }
    /*skip lines that begin with # */
    while (fgets(line, sizeof(line), file)) {
        if (line[0] == '#') {
            // Skip comment lines
            continue;
        }

        // Resize DataPoint data struct to hold more rows
        if (*row_count >= capacity) {
            capacity *= 2;
            *mag_anomaly = (DataPoint *)realloc(*mag_anomaly, capacity * sizeof(DataPoint));
       

            if (*mag_anomaly == NULL) {
               perror("Error allocating memory");
               fclose(file);
               return -1;
            }
        }

        /* Parse the line for the observed magnetic data
           using northing, easting, observed mag and calculated mag values
           separated by blank space.
           
           Only scan lines with four columns and track the row count 
        */
        if (sscanf(line, "%lf %lf %lf %lf", &(*mag_anomaly)[*row_count].northing,
                                        &(*mag_anomaly)[*row_count].easting,
                                        &(*mag_anomaly)[*row_count].mag_obs,
                                        &(*mag_anomaly)[*row_count].mag_calc) == 4) {
            (*row_count)++;
        }
    }

    fclose(file);
    return 0;
}