#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <check.h>
#if HAVE_CONFIG_H
#  include <config.h>
#endif

#include <freesasa.h>
#include <freesasa_internal.h>
#include "tools.h"

#define PASS 1
#define NOPASS 0

double total_ref, polar_ref, apolar_ref;
double tolerance;

freesasa_parameters parameters;

double rel_err(double v1, double v2) {
    return fabs(v1-v2)/(fabs(v1)+fabs(v2));
}

double surface_hidden_sphere_intersection(double r1, double r2, double d)
{
    if (d > r1 + r2) return 0;
    if (r1+d < r2) return 4*M_PI*r1*r1;
    if (r2+d < r1) return 4*M_PI*r2*r2;
    return M_PI/d * ( r1*(r2*r2 - (d-r1)*(d-r1))
                    +r2*(r1*r1 - (d-r2)*(d-r2)) ) ;
}
double surface_spheres_intersecting(double r1, double r2, double d)
{
    return 4*M_PI*(r1*r1+r2*r2) - surface_hidden_sphere_intersection(r1,r2,d);
}

double surface_two_spheres(const double *x, const double *r, double probe)
{
    double d2 = (x[0]-x[3])*(x[0]-x[3]) + (x[1]-x[4])*(x[1]-x[4])
        + (x[2]-x[5])*(x[2]-x[5]);
    return surface_spheres_intersecting(r[0]+probe,r[1]+probe,sqrt(d2));
}

int test_sasa(double ref, const char *test, const double *xyz, 
              const double *r, int n)
{
    double err;
    freesasa_result *result;
    int pass = PASS;
    result = freesasa_calc_coord(xyz,r,n,&parameters);
    if ((err = rel_err(ref,result->total))
        > tolerance) {
        pass = NOPASS;
    }
    freesasa_result_free(result);
    return pass;
}

void setup_lr_precision(void)
{
    parameters = freesasa_default_parameters;
    parameters.alg = FREESASA_LEE_RICHARDS;
    parameters.lee_richards_n_slices = 20000;
    tolerance = 1e-5;
}
void teardown_lr_precision(void)
{
    
}
void setup_sr_precision(void)
{
    parameters = freesasa_default_parameters;
    parameters.alg = FREESASA_SHRAKE_RUPLEY;
    parameters.shrake_rupley_n_points = 5000;
    tolerance = 1e-3;
}
void teardown_sr_precision(void)
{
    
}

START_TEST (test_sasa_alg_basic)
{
    // Two spheres, compare with analytic results
    double coord[6] = {0,0,0,2,0,0};
    double r[2] = {1,2};
    double probe = parameters.probe_radius;
    double n = 2;

    ck_assert(test_sasa(surface_two_spheres(coord,r,probe),
                        "Two intersecting spheres along x-axis.",
                        coord,r,n));
    coord[3] = 0; coord[4] = 2;
    ck_assert(test_sasa(surface_two_spheres(coord,r,probe),
                        "Two intersecting spheres along y-axis.",
                        coord,r,n));

    coord[4] = 0; coord[5] = 2;
    ck_assert(test_sasa(surface_two_spheres(coord,r,probe),
                        "Two intersecting spheres along z-axis.",
                        coord,r,n));

    // Four spheres in a plane, all calculations should give similar results
    double coord2[12] = {0,0,0, 1,0,0, 0,1,0, 1,1,0};
    double r2[4] = {1,1,2,1};
    n = 4;
    freesasa_result *result = freesasa_calc_coord(coord2,r2,4,&parameters);
    double ref = result->total;
    freesasa_result_free(result);

    //translate
    for (int i = 0; i < 12; ++i) coord2[i] += 1.;
    ck_assert(test_sasa(ref,"Four spheres in plane, translated",
                        coord2,r2,n));

    //rotate 90 degrees round z-axis
    double coord3[12] = {0,1,0, 0,0,0, 1,1,0, 1,0,0};
    memcpy(coord2,coord3,12*sizeof(double));
    ck_assert(test_sasa(ref,"Four spheres in plane, rotated 90 deg round z-axis.",
                        coord2,r2,n));

    //rotate -45 degrees round z-axis
    double sqr2 = sqrt(2);
    double coord4[12] = {-1./sqr2,1./sqr2,0, 0,0,0, 0,sqr2,0, 1/sqr2,1/sqr2,0};
    memcpy(coord2,coord4,12*sizeof(double));
    ck_assert(test_sasa(ref,"Four spheres in plane, rotated 45 deg round z-axis.",
                        coord2,r2,n));

    //rotate 90 degrees round x-axis
    double coord5[12] = {-1./sqr2,0,1/sqr2, 0,0,0, 0,0,sqr2, 1/sqr2,0,1/sqr2};
    memcpy(coord2,coord5,12*sizeof(double));
    ck_assert(test_sasa(ref,"Four spheres in plane, rotated 90 deg round x-axis.",
                        coord2,r2,n));
    
}
END_TEST

START_TEST (test_minimal_calc)
{
    double coord[3] = {0,0,0};
    double r[1] = {1.0};

    freesasa_result *result = freesasa_calc_coord(coord, r, 1, NULL);

    ck_assert_ptr_ne(result, NULL);

    // access areas
    ck_assert(fabs(result->sasa[0] - result->total) < 1e-10);
    ck_assert(fabs(result->total - (4*M_PI*M_PI*2.4*2.4)));
    
    freesasa_result_free(result);
}
END_TEST


void setup_sr (void)
{
    parameters = freesasa_default_parameters;
    parameters.alg = FREESASA_SHRAKE_RUPLEY;
    parameters.shrake_rupley_n_points = 100;
    parameters.n_threads = 1;
    total_ref = 4834.716265;
    polar_ref = 2515.821238;
    apolar_ref = 2318.895027;
}
void teardown_sr(void)
{

}

void setup_lr (void)
{
    parameters = freesasa_default_parameters;
    parameters.alg = FREESASA_LEE_RICHARDS;
    parameters.lee_richards_n_slices = 20;
    parameters.n_threads = 1;
    total_ref = 4804.055641;
    polar_ref = 2504.217302;
    apolar_ref = 2299.838339;
}
void teardown_lr(void)
{

}

START_TEST (test_sasa_1ubq)
{
    FILE *pdb = fopen(DATADIR "1ubq.pdb","r");
    ck_assert(pdb != NULL);
    const freesasa_classifier *classifier = &freesasa_default_classifier;
    freesasa_structure *st = freesasa_structure_from_pdb(pdb, classifier, 0);
    freesasa_result *res;
    freesasa_node *tree = freesasa_tree_new();
    ck_assert((res = freesasa_calc_structure(st, &parameters)) != NULL);

    freesasa_nodearea res_class = freesasa_result_classes(st, res);

    freesasa_tree_add_result(tree, res, st, "test");
    fclose(pdb);

    ck_assert(float_eq(res->total, total_ref, 1e-5));
    ck_assert(float_eq(res_class.polar, polar_ref, 1e-5));
    ck_assert(float_eq(res_class.apolar, apolar_ref, 1e-5));
    
    FILE *devnull = fopen("/dev/null","w");
    ck_assert(freesasa_write_pdb(devnull, tree) == FREESASA_SUCCESS);
    ck_assert(freesasa_tree_export(devnull, tree, FREESASA_PDB) == FREESASA_SUCCESS);
    ck_assert(freesasa_tree_export(devnull, tree, FREESASA_LOG) == FREESASA_SUCCESS);
    ck_assert(freesasa_tree_export(devnull, tree, FREESASA_RES) == FREESASA_SUCCESS);
    ck_assert(freesasa_tree_export(devnull, tree, FREESASA_SEQ) == FREESASA_SUCCESS);
    ck_assert(freesasa_tree_export(devnull, tree, FREESASA_RSA) == FREESASA_SUCCESS);
    if (USE_JSON) {
        ck_assert(freesasa_tree_export(devnull, tree, FREESASA_JSON) == FREESASA_SUCCESS);
    } else {
        ck_assert(freesasa_tree_export(devnull, tree, FREESASA_JSON) == FREESASA_FAIL);
    }
    if (USE_XML) {
        ck_assert(freesasa_tree_export(devnull, tree, FREESASA_XML) == FREESASA_SUCCESS);
    } else {
        ck_assert(freesasa_tree_export(devnull, tree, FREESASA_XML) == FREESASA_FAIL);
    }
    fclose(devnull);

    freesasa_set_verbosity(FREESASA_V_SILENT);
    FILE *nowrite = fopen("/dev/null","r");
    ck_assert(freesasa_write_log(nowrite, tree) == FREESASA_FAIL);
    ck_assert(freesasa_tree_export(nowrite, tree, FREESASA_RSA) == FREESASA_FAIL);
    ck_assert(freesasa_tree_export(nowrite, tree, FREESASA_JSON) == FREESASA_FAIL);
    if (USE_JSON) {
        ck_assert(freesasa_tree_export(nowrite, tree, FREESASA_JSON) == FREESASA_FAIL);
    }
    if (USE_XML) {
        ck_assert(freesasa_tree_export(nowrite, tree, FREESASA_XML) == FREESASA_FAIL);
    }
    fclose(nowrite);
    freesasa_set_verbosity(FREESASA_V_NORMAL);

    freesasa_structure_free(st);
    freesasa_result_free(res);
    freesasa_node_free(tree);
}
END_TEST

START_TEST (test_write_pdb) {
    FILE *tf = fopen("tmp/dummy_bfactors.pdb","w+"),
        *ref = fopen(DATADIR "reference_bfactors.pdb","r"),
        *pdb = fopen(DATADIR "1ubq.pdb","r"),
        *devnull = fopen("/dev/null", "w");
    ck_assert(tf != NULL);
    ck_assert(ref != NULL);
    ck_assert(pdb != NULL);
    freesasa_structure *s = freesasa_structure_from_pdb(pdb, NULL, 0);
    const int n = freesasa_structure_n(s);
    freesasa_result res;
    freesasa_node *root;
    fclose(pdb);

    res.sasa = malloc(sizeof(double)*n);
    for (int i = 0; i < n; ++i) res.sasa[i] = 1.23;
    res.parameters = freesasa_default_parameters;
    res.n_atoms = n;
    
    freesasa_structure_set_radius(s, res.sasa);
    root = freesasa_tree_init(&res, s, "bla");
    ck_assert(freesasa_write_pdb(tf, root) == FREESASA_SUCCESS);

    rewind(tf);
    free(res.sasa);

    //check that output matches reference file
    size_t bufsize = 100;
    char *buf_tf = malloc(bufsize), *buf_ref = malloc(bufsize);
    while(getline(&buf_tf,&bufsize,tf) > 0 && getline(&buf_ref,&bufsize,ref) > 0) {
        // skip remarks
        while (strncmp(buf_ref, "REMARK", 6) == 0 && getline(&buf_ref, &bufsize, ref) > 0);
        while (strncmp(buf_tf, "REMARK", 6) == 0 && getline(&buf_tf, &bufsize, tf) > 0);
        ck_assert_str_eq(buf_ref,buf_tf);
    }
    free(buf_tf);
    free(buf_ref);

    freesasa_structure_free(s);
    freesasa_node_free(root);

    // Can't write pdb from structure not initialized from pdb
    s = freesasa_structure_new();
    freesasa_structure_add_atom(s,"C","ALA","   1",'C',0,0,0);
    freesasa_set_verbosity(FREESASA_V_SILENT);
    root = freesasa_tree_init(&res, s, "bla");
    ck_assert_int_eq(freesasa_write_pdb(devnull, root), FREESASA_FAIL);
    freesasa_set_verbosity(FREESASA_V_NORMAL);

    freesasa_structure_free(s);
    freesasa_node_free(root);
    fclose(devnull);
    fclose(ref);
    fclose(tf);

}
END_TEST


START_TEST (test_trimmed_pdb) 
{
    // This test is due to suggestion from João Rodrigues (issue #6 on Github)
    double total_ref = 16133.867124;
    double polar_ref = 7432.608118;
    double apolar_ref = 8701.259006;
    freesasa_parameters param = freesasa_default_parameters;
    const freesasa_classifier *classifier = &freesasa_default_classifier;
    freesasa_result *result;
    freesasa_structure *st;
    freesasa_nodearea res_class;
    FILE *pdb;
    param.alg = FREESASA_SHRAKE_RUPLEY;

    pdb = fopen(DATADIR "3bzd_trimmed.pdb","r");
    ck_assert(pdb != NULL);
    st = freesasa_structure_from_pdb(pdb, classifier, 0);
    fclose(pdb);

    result = freesasa_calc_structure(st,&param);
    ck_assert_ptr_ne(result, NULL);
    res_class = freesasa_result_classes(st, result);

    ck_assert(float_eq(result->total, total_ref, 1e-5));
    ck_assert(float_eq(res_class.polar, polar_ref, 1e-5));
    ck_assert(float_eq(res_class.apolar, apolar_ref, 1e-5));

    freesasa_structure_free(st);
    freesasa_result_free(result);
}
END_TEST

START_TEST (test_user_classes)
{
    FILE *pdb = fopen(DATADIR "1ubq.pdb","r"),
        *clf = fopen(SHAREDIR "protor.config", "r");
    freesasa_structure *st, *st_ref;
    freesasa_classifier *user_classifier;
    freesasa_result *res, *res_ref;
    freesasa_nodearea res_class, res_class_ref;
    const double *radii, *radii_ref;

    ck_assert(pdb != NULL);
    ck_assert(clf != NULL);

    user_classifier = freesasa_classifier_from_file(clf);
    fclose(clf);
    ck_assert(user_classifier != NULL);

    st = freesasa_structure_from_pdb(pdb, user_classifier, 0);
    ck_assert(st != NULL);
    rewind(pdb);
    st_ref = freesasa_structure_from_pdb(pdb, &freesasa_protor_classifier, 0);
    ck_assert(st_ref != NULL);
    fclose(pdb);

    radii = freesasa_structure_radius(st);
    radii_ref = freesasa_structure_radius(st_ref);
    ck_assert(radii != NULL);
    ck_assert(radii_ref != NULL);
    for (int i = 0; i < freesasa_structure_n(st); ++i) {
        ck_assert(float_eq(radii[i], radii_ref[i], 1e-10));
    }
    res = freesasa_calc_structure(st, NULL);
    res_ref = freesasa_calc_structure(st_ref, NULL);
    ck_assert_ptr_ne(res, NULL);
    ck_assert_ptr_ne(res, NULL);
    res_class = freesasa_result_classes(st, res);
    res_class_ref = freesasa_result_classes(st_ref, res_ref);
    ck_assert(float_eq(res_class.total, res_class_ref.total, 1e-10));
    ck_assert(float_eq(res_class.polar, res_class_ref.polar, 1e-10));
    ck_assert(float_eq(res_class.apolar, res_class_ref.apolar, 1e-10));
    ck_assert(float_eq(res_class.unknown, res_class_ref.unknown, 1e-10));
    ck_assert(float_eq(res_class.side_chain, res_class_ref.side_chain, 1e-10));
    ck_assert(float_eq(res_class.main_chain, res_class_ref.main_chain, 1e-10));

    freesasa_structure_free(st);
    freesasa_classifier_free(user_classifier);
    freesasa_result_free(res);
}
END_TEST


START_TEST (test_calc_errors)
{
    freesasa_set_verbosity(FREESASA_V_NORMAL);

    fputs("Testing error messages:\n",stderr);
    freesasa_fail("Test fail-message.");
    freesasa_warn("Test warn-message.");
    ck_assert(freesasa_get_verbosity() == 0);
    freesasa_set_verbosity(FREESASA_V_SILENT);

    //test empty PDB-file
    FILE *empty = fopen(DATADIR "empty.pdb","r");
    ck_assert(empty != NULL);
    ck_assert(freesasa_structure_from_pdb(empty, NULL, 0) == NULL);
    fclose(empty);
    
    freesasa_set_verbosity(FREESASA_V_NORMAL);

}
END_TEST

START_TEST (test_multi_calc)
{
#if USE_THREADS
    FILE *pdb = fopen(DATADIR "1ubq.pdb","r");
    freesasa_structure *st = freesasa_structure_from_pdb(pdb, NULL, 0);
    freesasa_result *res;
    freesasa_parameters p = freesasa_default_parameters;

    fclose(pdb);

    //S&R
    p.n_threads = 2;
    p.alg = FREESASA_SHRAKE_RUPLEY;
    ck_assert((res = freesasa_calc_structure(st,&p)) != NULL);
    ck_assert(fabs(res->total - 4834.716265) < 1e-5);
    // L&R
    p.alg = FREESASA_LEE_RICHARDS;
    p.lee_richards_n_slices = 20;
    ck_assert((res = freesasa_calc_structure(st,&p)) != NULL);
    ck_assert(fabs(res->total - 4804.055641) < 1e-5);
    
    freesasa_structure_free(st);
    freesasa_result_free(res);
#endif /* USE_THREADS */
}
END_TEST

// test an NMR structure with hydrogens and several models
START_TEST (test_1d3z) 
{
    FILE *pdb = fopen(DATADIR "1d3z.pdb","r");
    int n = 0;
    freesasa_parameters param = freesasa_default_parameters;
    param.alg = FREESASA_SHRAKE_RUPLEY;
    freesasa_structure* st = freesasa_structure_from_pdb(pdb, NULL, 0);
    freesasa_result *result = freesasa_calc_structure(st, &param);
    double *radii_ref = malloc(sizeof(double)*602);
    ck_assert(freesasa_structure_n(st) == 602);
    ck_assert(fabs(result->total - 5000.340175) < 1e-5);
    memcpy(radii_ref,freesasa_structure_radius(st),sizeof(double)*602);
    rewind(pdb);
    freesasa_structure_free(st);
    
    freesasa_set_verbosity(FREESASA_V_SILENT);
    st = freesasa_structure_from_pdb(pdb, NULL, FREESASA_INCLUDE_HYDROGEN);
    result = freesasa_calc_structure(st, &param);
    ck_assert(freesasa_structure_n(st) == 1231);
    ck_assert(fabs(result->total - 5035.614493) < 1e-5);
    rewind(pdb);
    freesasa_set_verbosity(FREESASA_V_NORMAL);

    freesasa_structure** ss = freesasa_structure_array(pdb, &n, NULL, FREESASA_SEPARATE_MODELS);
    ck_assert(n == 10);
    result = freesasa_calc_structure(ss[0], &param);
    ck_assert(freesasa_structure_n(ss[0]) == 602);
    ck_assert(fabs(result->total - 5000.340175) < 1e-5);
    for (int i = 0; i < n; ++i) {
        const double *r2 = freesasa_structure_radius(ss[i]);
        ck_assert(r2 != NULL);
        for (int j = 0; j < 602; ++j) {
            ck_assert(fabs(r2[j] - radii_ref[j]) < 1e-10);
        }
        freesasa_structure_free(ss[i]);
    }
    free(ss);
    freesasa_result_free(result);
    freesasa_structure_free(st);
    fclose(pdb);
}
END_TEST

START_TEST (test_memerr)
{
    freesasa_parameters p = freesasa_default_parameters;
    double v[18] = {0,0,0, 1,1,1, -1,1,-1, 2,0,-2, 2,2,0, -5,5,5};
    struct coord_t coord = {.xyz = v, .n = 6, .is_linked = 0};
    const double r[6]  = {4,2,2,2,2,2};
    void *ptr;
    p.shrake_rupley_n_points = 10; // so the loop below will be fast

    freesasa_set_verbosity(FREESASA_V_SILENT);
    for (int i = 1; i < 40; ++i) {
        p.alg = FREESASA_SHRAKE_RUPLEY;
        set_fail_after(i);
        ptr = freesasa_calc(&coord, r, &p);
        set_fail_after(0);
        ck_assert_ptr_eq(ptr, NULL);
        p.alg = FREESASA_LEE_RICHARDS;
        set_fail_after(i);
        ptr = freesasa_calc(&coord, r, &p);
        set_fail_after(0);
        ck_assert_ptr_eq(ptr, NULL);
    }

    FILE *file = fopen(DATADIR "1ubq.pdb","r");
    freesasa_structure *s = freesasa_structure_from_pdb(file, NULL, 0);
    for (int i = 1; i < 256; i *= 2) { //try to spread it out without doing too many calculations
        set_fail_after(i);
        ptr = freesasa_calc_structure(s, NULL);
        set_fail_after(0);
        ck_assert_ptr_eq(ptr, NULL);
        set_fail_after(i);
        ptr = freesasa_structure_get_chains(s, "A");
        set_fail_after(0);
        ck_assert_ptr_eq(ptr, NULL);
    }
    freesasa_structure_free(s);
    fclose(file);
    freesasa_set_verbosity(FREESASA_V_NORMAL);
}
END_TEST

extern TCase * test_LR_static();

Suite *sasa_suite()
{
    Suite *s = suite_create("SASA-calculation");

    TCase *tc_basic = tcase_create("API");
    tcase_add_test(tc_basic, test_minimal_calc);
    tcase_add_test(tc_basic, test_calc_errors);
    tcase_add_test(tc_basic, test_user_classes);
    tcase_add_test(tc_basic, test_write_pdb);
    tcase_add_test(tc_basic, test_memerr);
    
    TCase *tc_lr_basic = tcase_create("Basic L&R");
    tcase_add_checked_fixture(tc_lr_basic,setup_lr_precision,teardown_lr_precision);
    tcase_add_test(tc_lr_basic, test_sasa_alg_basic);

    TCase *tc_lr_static = test_LR_static();
    
    TCase *tc_sr_basic = tcase_create("Basic S&R");
    tcase_add_checked_fixture(tc_sr_basic,setup_sr_precision,teardown_sr_precision);
    tcase_add_test(tc_sr_basic, test_sasa_alg_basic);

    TCase *tc_lr = tcase_create("1UBQ-L&R");
    tcase_add_checked_fixture(tc_lr,setup_lr,teardown_lr);
    tcase_add_test(tc_lr, test_sasa_1ubq);

    TCase *tc_sr = tcase_create("1UBQ-S&R");
    tcase_add_checked_fixture(tc_sr,setup_sr,teardown_sr);
    tcase_add_test(tc_sr, test_sasa_1ubq);

    TCase *tc_trimmed = tcase_create("Trimmed PDB file");
    tcase_add_test(tc_trimmed, test_trimmed_pdb);

    TCase *tc_1d3z = tcase_create("NMR PDB-file 1D3Z (several models, hydrogens)");
    tcase_add_test(tc_1d3z,test_1d3z);

    suite_add_tcase(s, tc_basic);
    suite_add_tcase(s, tc_lr_basic);
    suite_add_tcase(s, tc_lr_static);
    suite_add_tcase(s, tc_sr_basic);
    suite_add_tcase(s, tc_lr);
    suite_add_tcase(s, tc_sr);
    suite_add_tcase(s, tc_trimmed);
    suite_add_tcase(s, tc_1d3z);

#if USE_THREADS
    printf("Using pthread\n");
    TCase *tc_pthr = tcase_create("Pthread");
    tcase_add_test(tc_pthr,test_multi_calc);
    suite_add_tcase(s, tc_pthr);
#endif
    return s;
}
