import numpy as np

cimport numpy as np
#from libc.stdlib cimport free

cdef extern from "tng/tng_io_fwd.h":
    cdef struct tng_trajectory:
        pass

    ctypedef tng_trajectory *tng_trajectory_t


cdef extern from "tng/tng_io.h":
    ctypedef int int64_t
    cdef int64_t TNG_TRAJ_POSITIONS "TNG_TRAJ_POSITIONS"
    cdef int64_t TNG_TRAJ_VELOCITIES "TNG_TRAJ_VELOCITIES"
    cdef int64_t TNG_TRAJ_FORCES "TNG_TRAJ_FORCES"

    ctypedef enum tng_function_status:
         TNG_SUCCESS, TNG_FAILURE, TNG_CRITICAL

    tng_function_status tng_util_trajectory_open(const char *filename, const char mode, tng_trajectory_t *tng_data_p)

    tng_function_status tng_util_trajectory_close(tng_trajectory_t *tng_data_p)

    tng_function_status tng_util_time_of_frame_get(const tng_trajectory_t tng_data, const int64_t frame_nr, double *time)

    tng_function_status tng_util_pos_read_range(const tng_trajectory_t tng_data,
                                                const int64_t first_frame,
                                                const int64_t last_frame,
                                                float **positions,
                                                int64_t *stride_length)

    tng_function_status tng_util_vel_read_range(const tng_trajectory_t tng_data,
                                                const int64_t first_frame,
                                                const int64_t last_frame,
                                                float **velocities,
                                                int64_t *stride_length)

    tng_function_status tng_util_force_read_range(const tng_trajectory_t tng_data,
                                                  const int64_t first_frame,
                                                  const int64_t last_frame,
                                                  float **forces,
                                                  int64_t *stride_length)

    tng_function_status tng_util_box_shape_read_range(const tng_trajectory_t tng_data,
                                                      const int64_t first_frame,
                                                      const int64_t last_frame,
                                                      float **box_shape,
                                                      int64_t *stride_length)

    tng_function_status tng_num_frames_get(const tng_trajectory_t tng_data, int64_t *n)

    tng_function_status tng_num_particles_get(const tng_trajectory_t tng_data, int64_t *n);

    tng_function_status tng_medium_stride_length_get(const tng_trajectory_t tng_data, int64_t *len);

    tng_function_status tng_long_stride_length_get(const tng_trajectory_t tng_data, int64_t *len);

    tng_function_status tng_util_trajectory_next_frame_present_data_blocks_find(
                    const tng_trajectory_t tng_data,
                    int64_t current_frame,
                    const int64_t n_requested_data_block_ids,
                    const int64_t *requested_data_block_ids,
                    int64_t *next_frame,
                    int64_t *n_data_blocks_in_next_frame,
                    int64_t **data_block_ids_in_next_frame)

    tng_function_status tng_util_particle_data_next_frame_read(
                    const tng_trajectory_t tng_data,
                    const int64_t block_id,
                    void **values,
                    char *data_type,
                    int64_t *retrieved_frame_number,
                    double *retrieved_time)


cdef check_status(status, message=''):
    if status != TNG_SUCCESS:
        print(message)

#cdef ptr_to_array(float *ptr, np.npy_intp *shape, int np):
#    arr = np.PyArray_SimpleNewFromData(np, shape, np.NPY_FLOAT, ptr)
#    return arr

cdef class TNGTrajectory:
    cdef:
        tng_trajectory_t trajectory
        int nframes, natoms, stride

    def get_position_interval(self):
        cdef:
            double time
            char data_type
            int frame1, frame2

        cdef np.ndarray[float, ndim=2] positions = np.empty((self.natoms, 3), dtype=np.float32)
        cdef void *data = <void *>positions.data

        tng_util_particle_data_next_frame_read(self.trajectory, TNG_TRAJ_POSITIONS, &data, &data_type, &frame1, &time)
        tng_util_particle_data_next_frame_read(self.trajectory, TNG_TRAJ_POSITIONS, &data, &data_type, &frame2, &time)
        del positions
        return frame2 - frame1

    def __len__(self):
        return self.nframes

    def __getitem__(self, frame):
        cdef np.ndarray[float, ndim=2] positions = np.empty((self.natoms, 3), dtype=np.float32)
        cdef float *pos = <float *>positions.data

        #check_status(
        if (TNG_SUCCESS !=
            tng_util_pos_read_range(self.trajectory, frame, frame, &pos, &self.stride)#,
            #'Reading of frame not successful.'
            ):
            raise IndexError()

        return positions

    def __init__(self, filename):
        assert filename, 'Filename must not be null.'
        if isinstance(filename, str):
            filename = filename.encode()

        if (tng_util_trajectory_open(filename, b'r', &self.trajectory) != TNG_SUCCESS):
            raise OSError('Trajectory could not be opened.')

        tng_num_frames_get(self.trajectory, &self.nframes)
        tng_num_particles_get(self.trajectory, &self.natoms)


from math import ceil

class TNGReader:
    def __init__(self, filename):
        self.trajectory = TNGTrajectory(filename)
        self.nstxout = self.trajectory.get_position_interval()

    def __len__(self):
        return ceil(len(self.trajectory) / self.nstxout)

    def __getitem__(self, item):
        try:
            return self.trajectory[item * self.nstxout]
        except IndexError:
            raise IndexError('Trajectory frame {} out of range.'.format(item))