#!/usr/bin/env python
import numpy as np
import scipy
import Fold_utils


#------------------------------------------------------------------------------
## This script generates a text file containing fake data from an FRB


#------------------------------------------------------------------------------
## Observing setup

## Input parameters for the observing setup
nchannels = 512
bottom_freq = 1200.
bandwidth = 400.
dt_sampling = 5e-5
obs_length = 1.

## Derived quantities
chanwidth = bandwidth / nchannels
freqs = np.arange(bottom_freq, bottom_freq+bandwidth, chanwidth)
tstamp = np.arange(0., obs_length, dt_sampling)


#------------------------------------------------------------------------------
## Fake burst setup

## Fake burst properties
t_burst = 0.1
dm_burst = 150.
duration_burst = 2.65e-3
amp_burst = 0.5

## Generating the model data for the fake burst
t_burst_at_freqs = Fold_utils.DM_delay(dm_burst, freqs)
model_data = np.exp(-0.5*(tstamp[:,None]-(t_burst+t_burst_at_freqs[None,:]))**2/duration_burst**2)

## Preparing the file header
header = (
    '# number of channels: {}\n'
    '# bottom frequency (MHz): {}\n'
    '# bandwidth (MHz): {}\n'
    '# sampling time (s): {}\n'
    '# file format: columns -> frequencies (increasing), rows -> time (increasing)'
    ).format(nchannels, bottom_freq, bandwidth, dt_sampling)


#------------------------------------------------------------------------------
## Generating data

## Dumping the model data to a file
np.savetxt('frb_model.txt', model_data, header=header, comments='', fmt='%.4e')

## Adding gaussian noise
## Each channel has its own noise, drawn from a lognormal distribution
noise_std_per_channel = np.random.lognormal(mean=0.0, sigma=0.1, size=nchannels)
noise = np.random.normal(loc=0., scale=noise_std_per_channel, size=model_data.shape)
np.savetxt('frb_noise.txt', model_data+noise, header=header, comments='', fmt='%.4e')