import numpy as np

class Vertex(object):
    def __init__(self, particle_name, pos, dir, ke, t0=0.0, pol=None):
        self.particle_name = particle_name
        self.pos = pos
        self.dir = dir
        self.pol = pol
        self.ke = ke
        self.t0 = t0

class Photons(object):
    def __init__(self, pos, dir, pol, wavelengths, t=None, last_hit_triangles=None, flags=None):
        self.pos = pos
        self.dir = dir
        self.pol = pol
        self.wavelengths = wavelengths

        if t is None:
            self.t = np.zeros(len(pos), dtype=np.float32)
        else:
            self.t = t

        if last_hit_triangles is None:
            self.last_hit_triangles = np.empty(len(pos), dtype=np.int32)
            self.last_hit_triangles.fill(-1)
        else:
            self.last_hit_triangles = last_hit_triangles

        if flags is None:
            self.flags = np.zeros(len(pos), dtype=np.uint32)
        else:
            self.flags = flags

    def __add__(self, other
/* Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.

 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.

 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "zebra.h"
#include <stdio.h> /* for FILE */
#include <stdlib.h> /* for size_t */
#include "pack2b.h"
#include <string.h> /* for memmove() */
#include <stdint.h> /* for uint8_t, etc. */
#include <errno.h> /* for strerror(), etc. */
#include <zlib.h>

char zebra_err[256];

/* Returns a pointer to a new zebra file object.
 *
 * Returns the pointer to the zebraFile object on success, or NULL on error. */
zebraFile *zebra_open(const char *filename)
{
    /* We use gzopen() here even though the file may not be gzipped. According
     * to the docs:
     *
     *      gzopen can be used to read a file which is not in gzip format; in
     *      this case gzread will directly read from the file without
     *      decompression.  When reading, this will be detected automatically
     *      by looking for the magic two-byte gzip header.
     *
     * See https://zlib.net/manual.html. */
    gzFile f = gzopen(filename, "r");

    if (!f) {
        sprintf(zebra_err, "failed to open '%s': %s", filename, strerror(errno));
        return NULL;
    }

    zebraFile *z = (zebraFile *) malloc(sizeof(zebraFile));

    if (!z) {
        sprintf(zebra_err, "failed to allocate space for zebraFile struct: %s", strerror(errno));
        gzclose(f);
        return NULL;
    }

    z->f = f;
    z->lr_size = 0;
    z->buf = NULL;
    z->buf_size = 0;
    z->tab = NULL;

    return z;
}

static int read_next_physical_record(zebraFile *z)
{
    /* Read the next physical record from the zebra file and append it to the
     * buffer.
     *
     * Returns 0 on success, -1 on error, and 1 on EOF. */
    uint8_t buf[48];
    size_t nbytes;
    uint32_t s0, s1, s2, s3, nwphr, nfast;

    nbytes = gzread(z->f,buf,4*8);

    if (nbytes == 0) return 1;

    if (nbytes != 8*4) {
        sprintf(zebra_err, "expected %i bytes but only read %zu", 8*4, nbytes);
        return -1;
    }

    s0 = unpacki32(buf);
    s1 = unpacki32(buf+4*1);
    s2 = unpacki32(buf+4*2);
    s3 = unpacki32(buf+4*3);
    nwphr = unpacki32(buf+4*4);
    nfast = unpacki32(buf+7*4);

    if (nwphr & (ZEBRA_EMERGENCY_STOP | ZEBRA_END_OF_RUN))
        return 1;

    nwphr &= ZEBRA_BLOCK_SIZE_MASK;

    uint32_t words = nwphr*(1 + nfast) - 8;

    if (s0 != ZEBRA_SIG0) {
        sprintf(zebra_err, "invalid steering block stamp 0x%08x", s0);
        return -1;
    }

    if (s1 != ZEBRA_SIG1) {
        sprintf(zebra_err, "invalid steering block stamp 0x%08x", s1);
        return -1;
    }

    if (s2 != ZEBRA_SIG2) {
        sprintf(zebra_err, "invalid steering block stamp 0x%08x", s2);
        return -1;
    }

    if (s3 != ZEBRA_SIG3) {
        sprintf(zebra_err, "invalid steering block stamp 0x%08x", s3);
        return -1;
    }

    z->buf = realloc(z->buf, z->buf_size+words*4);

    nbytes = gzread(z->f, z->buf+z->buf_size, 4*words);

    z->buf_size += words*4;

    if (nbytes != words*4) {
        sprintf(zebra_err, "expected %i bytes but only read %zu", words*4, nbytes);
        return -1;
    }

    return 0;
}

static int get_bytes(zebraFile *z, size_t size)
{
    /* Read bytes from the zdab file until there are at least `size` words in
     * the buffer.
     *
     * Returns 0 on success, -1 on error, and 1 on EOF. */
    int rv;

    while (z->buf_size < size*4) {
        rv = read_next_physical_record(z);

        if (rv == -1) {
            return -1;
        } else if (z->buf_size == 0 && rv == 1) {
            /* EOF and there are no bytes left. */
            return 1;
        } else if (rv == 1) {
            /* EOF but there are bytes left. */
            sprintf(zebra_err, "got EOF but there are still %zu bytes in the buffer!", z->buf_size);
            return -1;
        }
    }

    return 0;
}

static uint32_t get_link(zebraFile *z, uint32_t link)
{
    /* Returns the offset for a zebra bank in the local buffer from a link
     * pointing to the original bank. */
    int i;
    uint32_t sum;

    if (z->nwtab == 0) return link;

    sum = 0;
    for (i = 0; i < z->nwtab/2; i++) {
        if (link <= z->tab[i*2+1]) break;
        sum += z->tab[i*2+1] - z->tab[i*2];
    }

    return sum + link - z->tab[i*2] + z->lr_offset;
}

static int rewrite_links(zebraFile *z)
{
    /* Rewrite the next/up/orig, structural, and reference links for all banks
     * in the current logical record.
     *
     * Returns 0 on success, -1 on error. */
    int i, rv;
    uint32_t offset;
    uint32_t io, noff;
    zebraBank b;

    offset = z->lr_offset*4;

    z->mast_bank = -1;
    while (offset < z->lr_size) {
        io = unpacki32(z->buf+offset);
        offset += 4;

        noff = io & 0xFFFF;

        if (noff > 12) {
            offset += (noff - 12)*4;
        }

        rv = zebra_get_bank(z, &b, offset/4);

        if (!strncmp(b.name,"MAST",4)) {
            if (z->mast_bank != -1)
                fprintf(stderr, "Warning: found more than one MAST bank in a single logical record!\n");
            z->mast_bank = offset/4;
        }

        if (rv) return rv;

        /* Rewrite the next, up, and orig pointers. */
        if (b.next)
            packi32(z->buf+offset,get_link(z,b.next));
        if (b.up)
            packi32(z->buf+offset+4,get_link(z,b.up));
        if (b.orig)
            packi32(z->buf+offset+8,get_link(z,b.orig));

        /* Rewrite the structural and reference links. */
        for (i = 0; i < b.num_links; i++) {
            if (b.links[i])
                packi32(z->buf+offset-(i+1)*4,get_link(z,b.links[i]));
        }

        offset += 9*4;

        offset += b.num_data_words*4;
    }

    return 0;
}

int zebra_read_next_logical_record(zebraFile *z)
{
    /* Read the next logical record into the buffer.
     *
     * Returns 0 on success, -1 on error, and 1 on EOF. */
    int i;
    uint32_t size, type, nwtx, nwseg, nwtab, nwuh;
    uint32_t io, noff;
    uint32_t offset;
    int rv;

    while (1) {
        /* Move the next logical record to the start of the buffer. */
        memmove(z->buf,z->buf+z->lr_size,z->buf_size-z->lr_size);
        z->buf = realloc(z->buf, z->buf_size-z->lr_size);
        z->buf_size -= z->lr_size;
        offset = 0;

        switch (get_bytes(z,offset+1)) {
        case 1:
            return 1;
        case -1:
            return -1;
        }
        size = unpacki32(z->buf+offset*4);
        offset += 1;

        /* If the size of the logical record is zero, we just skip it.
         *
         * From page 142 of the ZEBRA document at
         * https://cdsweb.cern.ch/record/2296399/files/zebra.pdf:
         *
         * "The total size of a padding record is NWLR+1, including the LR
         * length and the LR type; thus a padding record with NWLR=1 occupies
         * two words. To pad exactly one word, NWLR=0 should be stored, in this
         * case the LR type 5 is implied without being present in the data." */
        if (size == 0) {
            z->lr_size = 4;
            continue;
        }

        switch (get_bytes(z,offset+1)) {
        case 1:
            return 1;
        case -1:
            return -1;
        }

        /* Get the logical record type. */
        type = unpacki32(z->buf+offset*4);
        offset += 1;

        switch (type) {
        /* Padding records. */
        case 5:
        case 6:
            if ((rv = get_bytes(z,offset+size-1))) return rv;
            z->lr_size = (size + 1)*4;
            continue;
        /* Start or end of run. */
        case 1:
            if ((rv = get_bytes(z,offset+size))) return rv;
            z->lr_size = (size + 2)*4;
            continue;
        /* Normal records. */
        case 2:
        case 3:
        case 4:
            break;
        }

        /* For normal records the size of the logical record is NWLR + the two
         * words for the size and type. */
        z->lr_size = (size + 2)*4;
        if ((rv = get_bytes(z,offset+size))) return rv;

        nwtx = unpacki32(z->buf+(offset+4)*4);
        nwseg = unpacki32(z->buf+(offset+5)*4);
        nwtab = unpacki32(z->buf+(offset+6)*4);
        nwuh = unpacki32(z->buf+(offset+9)*4);

        if (nwtab % 2) {
            sprintf(zebra_err, "number of words in relocation table is not a multiple of 2!");
            return -1;
        }

        offset += 10;

        /* Currently we don't use the user header, segment table, or the text
         * vector. */
        offset += nwuh;
        offset += nwseg;
        offset += nwtx;

        /* Read the relocation table. */
        if (z->tab) free(z->tab);
        z->tab = malloc(sizeof(double)*nwtab);
        for (i = 0; i < nwtab; i++)
            z->tab[i] = unpacki32(z->buf+(offset+i)*4);
        z->nwtab = nwtab;
        offset += nwtab;

        z->lr_offset = offset;

        /* Find the pointer to the first bank. */
        io = unpacki32(z->buf+offset*4);

        noff = io & 0xFFFF;

        /* First bank is just past the end of the logical record (one word for
         * the io control word). */
        z->first_bank = offset + 1;

        /* Skip the links to get to the bank's next pointer. */
        if (noff > 12) {
            z->first_bank += noff - 12;
        }

        /* Rewrite the next/up/orig, structural, and reference links. */
        if (rewrite_links(z)) return -1;

        break;
    }

    return 0;
}

int zebra_get_bank(zebraFile *z, zebraBank *b, uint32_t link)
{
    /* Get the bank located at `link`.
     *
     * On success 0 is returned and the zebraBank struct `b` is filled with the
     * bank information. On error, -1 is returned and zebra_err is set. */
    int i;

    if ((link+9)*4 > z->buf_size) {
        sprintf(zebra_err, "link location of %i is outside buffer!", link);
        return -1;
    }

    unpack(z->buf+link*4,"lllllllll",&b->next, &b->up, &b->orig, &b->number, &b->idh, &b->num_links, &b->num_structural_links, &b->num_data_words, &b->status);

    unpack(z->buf+link*4+16,"cccc",b->name,b->name+1,b->name+2,b->name+3);

    b->name[4] = '\0';

    if (b->num_links > MAX_LINKS) {
        sprintf(zebra_err, "number of links in bank is greater than %i!", MAX_LINKS);
        return -1;
    }

    if (link < b->num_links) {
        sprintf(zebra_err, "buffer only has %i words before bank, not enough for %i links!", link, b->num_links);
        return -1;
    }

    /* Read in reference and structural links. */
    for (i = 0; i < b->num_links; i++)
        b->links[i] = unpacki32(z->buf+(link-1-i)*4);

    link += 9;

    if ((link+b->num_data_words)*4 > z->buf_size) {
        sprintf(zebra_err, "bank contains %i data bytes, but there are only %zu bytes left in buffer!", b->num_data_words*4, z->buf_size-link*4);
        return -1;
    }

    b->data = (uint32_t *) (z->buf+link*4);

    return 0;
}

void zebra_close(zebraFile *z)
{
    gzclose(z->f);
    if (z->tab) free(z->tab);
    if (z->buf) free(z->buf);
    free(z);
}