from __future__ import print_function, division, absolute_import
from builtins import range

import os
import sys
import numpy as np
import time
import matplotlib
# Force matplotlib to not use any Xwindows backend
matplotlib.use('Agg')
import matplotlib.pyplot as plt

########################################

def read_sos_rad(file_path, max_teta=-1):

    res = np.genfromtxt(file_path, names=['phi', 'teta','I','Q', 'U'])
    
    phis  = np.unique(res['phi'])
    tetas = np.unique(res['teta'])

    if max_teta!=-1:
        tetas = tetas[tetas<max_teta]

    rad = np.zeros(( len(tetas), len(phis), 3))

    for iphi,phi in enumerate(phis):
        for iteta,teta in enumerate(tetas):
            #print(phi,teta)
            i = np.argwhere( (res['phi'] == phi) & (res['teta'] == teta) )
            #print(i, res['phi'][i], res['teta'][i])
            rad[iteta,iphi,0] = res['I'][i]
            rad[iteta,iphi,1] = res['Q'][i]
            rad[iteta,iphi,2] = res['U'][i]
            
    return phis, tetas, rad


if __name__=='__main__':
    
    phis, tetas, rad = read_sos_rad("/home/mathieu/work/RTM/RadiativeTransferCode-SOS/SOS_TEST/SOS/SOS_Up.txt")