#!/usr/bin/env python2
from __future__ import print_function

import numpy as np
import filterbank
import sigproc
from astropy.time import Time
from astropy import units as u

import os.path
import os
import sys
import re

import argparse

telescope_found = False
try:
    from telescope import telescope
    telescope_found = True
except:
    print("Could not connect to console, no RA/DEC recorded")

def create_filterbank(outfile, infile=None, date=None, source=None, create=False):
    """
    Create a filterbank file from the CAMRAS-format, which is generated
    by pulsar_record.

    outfile:  name of the filterbank file to be created
    infile:   name of the file generated by pulsar_record
    date:     the date to be used; default is the file creation date of infile
              if given, the date must be in a format the astropy.time understands
    source:   name of the source to be used; default is guessed from filename
    """

    if not source:
        if not infile:
            raise RuntimeError("If no input file is given, you must specify source")
        match = re.match("[BJ][0-9]*[+-][0-9]*", infile)
        if match:
            source = "PSR " + match.group(0)
        else:
            raise ValueError("Could not guess source name from filename")

    if not date:
        if infile:
            date = Time(os.path.getctime(infile), format='unix')
        else:
            date = Time.now()
    else:
        date = Time(date, format='isot')

    az_start = 0.
    za_start = 0.
    src_raj = 0.
    src_dej = 0.
    if telescope_found:
        dt = telescope(consoleHost='console')
        az_start = (180*u.deg + dt.getAzEl(waitForUpdate=True)[0]).to(u.deg).value
        za_start = (90*u.deg - dt.getAzEl()[1]).to(u.deg).value
        src_raj = dt.getRaDec().ra.to_string(unit=u.hour, sep='')
        src_dej = dt.getRaDec().dec.to_string(unit=u.degree, sep='')

    print("Input file: {}".format(infile))
    print("Output file: {}".format(outfile))
    print("Source name: {}".format(source))
    print("Observation time: {}".format(date.isot))

    date = date.mjd

    if infile:
        rawdatafile = infile
    else:
        rawdatafile = ""
 
    # Generate filterbank header
    fil_header = {}
    fil_header["telescope_id"]  = sigproc.telescope_ids["Effelsberg"]
    fil_header["machine_id"]    = sigproc.machine_ids["FAKE"]
    fil_header["data_type"]     = 1
    fil_header["rawdatafile"]   = rawdatafile
    fil_header["source_name"]   = source
    fil_header["barycentric"]   = 0
    fil_header["pulsarcentric"] = 0
    fil_header["az_start"]      = az_start
    fil_header["za_start"]      = za_start
    fil_header["src_raj"]       = src_raj
    fil_header["src_dej"]       = src_dej
    fil_header["tstart"]        = date
    fil_header["tsamp"]         = 0.00046811428571414528 # 1.0/35e6*256 chans*64 decimation
    fil_header["nbits"]         = 16
    fil_header["fch1"]          = 441.4
    fil_header["foff"]          = -35.0/256.0
    fil_header["nchans"]        = 256
    fil_header["nifs"]          = 1

    # Write header
    out = filterbank.create_filterbank_file(outfile, fil_header, nbits=16)

    if create:
        os.system("./pulsar_filterbank 3600 "+outfile)
        return

    if not infile:
        out.close()
        return

    # Read file
    f = open(infile)

    filesize = os.stat(infile).st_size
    offset = 0
    CHUNKSIZE = 30000
    while offset < filesize:
        # Loop through file in chunks of 256*10000 numbers
        # i.e. 256*CHUNKSIZE*4 bytes (10Mb) to avoid memory problems
        f.seek(offset, os.SEEK_SET)

        data = np.fliplr(np.fromfile(f, dtype='>u4', count=256*CHUNKSIZE).reshape(-1,256))

        # Check that no packets were dropped
        if np.sum(data[1:,-1] - data[:-1,-1] - 1) > 0.1:
            print("Packets were dropped:", np.count_nonzero(data[:-1,-1] - data[1:,-1] - 1),
                                           np.sum(data[1:,-1] - data[:-1,-1] - 1))

        # Write data
        out.append_spectra(data.astype('u2'))

        offset += 256*CHUNKSIZE*4

    # Close file
    out.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert a camras pulsar file (output of pulsar_record) to filterbank format")
    parser.add_argument("outfile", help="Name of output file, e.g. 'B0329.fil'")
    parser.add_argument("-i", "--infile", help="Input CAMRAS pulsar file. If empty, create just the header.")
    parser.add_argument("-d", "--date", help="Start date/time of observation, in 'isot' format (defaults to creation date of file, or now if no infile given)", default=None)
    parser.add_argument("-s", "--source", help="Name of the source, e.g. 'PSR B0329+54' (default: guessed from input filename)", default=None)
    parser.add_argument("-c", "--create", help="Create new observation by calling pulsar_filterbank", default=False, action="store_true")

    args = parser.parse_args()

    if args.infile is None and not(args.create):
        parser.error("One of --infile or --create is required")
    if args.infile is None and args.source is None:
        parser.error("Without infile, --source must be specified")
 
    create_filterbank(args.outfile, infile=args.infile, date=args.date, source=args.source, create=args.create)