import numpy as np
import math

from astropy import units as ut
from matplotlib import pyplot as plt
from matplotlib.offsetbox import AnchoredText
from os import path, makedirs
import base64
import io

# -- below are slight modifications of
# original author: Min-Su Shin , University of Michigan ----- #-#

def linear(inputArray, scale_min=None, scale_max=None):
    """
Performs linear scaling of the input numpy array.

  :inputArray:
    numpy array
    image data array

  :scale_min:
    float
    minimum data value
  :scale_max:
    float
    maximum data value
  :return:
    numpy array
    image data array

	"""
    if not inputArray.max==0:
      imageData=np.array(inputArray, copy=True)

      if scale_min == None:
          scale_min = imageData.min()
      if scale_max == None:
          scale_max = imageData.max()

      imageData = imageData.clip(min=scale_min, max=scale_max)
      imageData = (imageData -scale_min) / (scale_max - scale_min)
      indices = np.where(imageData < 0)
      imageData[indices] = 0.0
      indices = np.where(imageData > 1)
      imageData[indices] = 1.0

    return imageData

def sqrt(inputArray, scale_min=None, scale_max=None):
    """
Performs sqrt scaling of the input numpy array.

  :inputArray:
    numpy array
    image data array

  :scale_min:
    float
    minimum data value
  :scale_max:
    float
    maximum data value
  :return:
    numpy array
    image data array


	"""
    if not inputArray.max == 0:
        imageData=np.array(inputArray, copy=True)

    if scale_min == None:
        scale_min = imageData.min()
    if scale_max == None:
        scale_max = imageData.max()

    imageData = imageData.clip(min=scale_min, max=scale_max)
    imageData = imageData - scale_min
    indices = np.where(imageData < 0)
    imageData[indices] = 0.0
    imageData = np.sqrt(imageData)
    imageData = imageData / math.sqrt(scale_max - scale_min)

    return imageData


def log(inputArray, scale_min=None, scale_max=None,factor=2.0):
    """
Performs log10 scaling of the input numpy array.

  :inputArray:
    numpy array
    image data array

  :scale_min:
    float
    minimum data value

  :scale_max:
    float
    maximum data value

  :factor:
    float
    logarithm multiplication factor.

  :return:
    numpy array
    image data array
	"""
    if not inputArray.max == 0:
        imageData = np.array(inputArray, copy=True)


    if scale_min == None:
        scale_min = imageData.min()
    if scale_max == None:
        scale_max = imageData.max()
    #factor =math.log10(scale_max - scale_min)
    indices0 = np.where(imageData < 0)
    indices1 = np.where((imageData >= scale_min) & (imageData <= scale_max))
    indices2 = np.where(imageData > scale_max)
    imageData[indices0] = 0.0
    imageData[indices2] = 1.0
    try:
        imageData[indices1] = np.log10(imageData[indices1])/factor
    except:
        print ("Error on math.log10 ")
    return imageData

# ---------- author : @avialxee ---------#-#


def save_fig(plt, fig, kind='base64', output='output.jpg'):
    
    if kind == 'base64':
        buf = io.BytesIO()
        fig.savefig(buf, format='png', bbox_inches='tight',
                    transparent=True, pad_inches=0)
        buf.seek(0)
        string = base64.b64encode(buf.read())
        plt.close()
        return string
    elif kind == 'plot':
        plt.show()
        return 'plotted'
    else :
        if not path.exists('output'):
            makedirs('output')
        newPath = 'output/'+output
        opt = newPath
        if path.exists(newPath):
            numb = 1
            while path.exists(newPath):
                newPath = "{0}_{2}{1}".format(
                    *path.splitext(opt) + (numb,))
                try :
                    if path.exists(newPath):
                        numb += 1 
                except:
                    pass               
        fig.savefig(newPath, format=kind, bbox_inches='tight',
                    pad_inches=0)
        print("saved {}".format(newPath))
        plt.close()
        return newPath


def overlayc (r,g,b,c,lvl,cmin) :
  """
**Returns RGB stacked image and contour levels to be overlayed.**


*Input*

  :r/g/b:
    (2-D array)

    survey 2-D array input data

  :c:
    (2-D array)

    survey data for contour

  :lvl:
    (int)

    number of levels to draw

  :cmin:
    (float)
    minimum value selected for the contour

*Returns*

  :Img:
    (np array)
    dimension: (px,px,3)

  :lvlc:
    (list)
    dimension: None or of length len(lvl)
  """
  ri = normals(r)
  gi = normals(g)
  bi = normals(b)
  if ri.max() != 0 and ri.min()!=0 :
    ri = sqrt(ri,scale_min=0.1*np.std(ri),scale_max=np.max(ri))
  if gi.max() != 0 and gi.min()!=0 :
    gi = sqrt(gi, scale_min=0.1*np.std(gi), scale_max=np.max(gi))
  if bi.max() != 0 and bi.min()!=0 :
    bi = sqrt(bi,scale_min=0.1*np.std(bi),scale_max=np.max(bi))
  replace = bi > 1
  bi[replace] = 1
  replace = gi > 1
  gi[replace] = 1
  replace = ri > 1
  ri[replace] = 1

  # Stacking the Images and set up as 8-bit integers (highest value 2^8 = 256)
  img = (np.dstack((ri,gi,bi))*255.99).astype(np.uint8)
  if c.max() > cmin :
    lvlc = np.arange(cmin, c.max(),((c.max() - cmin)/lvl))
  else :
    lvlc = None
  return img, lvlc

def normals(o) :
  """
normalizing before scaling

*Input*

  :o:
    (2D array)
    input survey

*Returns*

  :(2D array):
    dimension same as input survey

  """
  X_scaled =(o - np.median(o))
  if o.max()!=0 :
    X = X_scaled/o.max()
    return X
  else :
    return X_scaled

def pl_powerlawsi(S,S_e,freq= [150, 1420], kind='plot', label="output"):
  """
plot powerlaw of spectralindex

*Inputs*

  :S:
    Total flux of TGSS and NVSS resp (list)
  :S_e:
    error on Total flux of TGSS and NVSS resp (list)
  :freq:
    frequency of TGSS and NVSS resp (list) (default = [150, 1420] )
  :kind:
    (str) choose from 'base64', 'plot', 'jpg','png'
      - 'base64' : base64 encoded
      - 'plot' : outputs plot in window or in your jupyter notebook
      - 'jpg' : saves jpg image output in output/
      - 'png' : saves png image output in output/

*Example*

.. code-block:: python

  >>> S_e = [42.6, 2.6]
  >>> S = [424.0, 60.0]
  >>> pl_si(S,S_e)
  """

  tgss = [float(S[0])]
  nvss = [float(S[1])]
  factor = freq[0]/freq[1]
  si = np.round(np.log(np.divide(tgss, nvss)+1E-5)/np.log(factor+1E-5), 3)
  fig = plt.figure(figsize=(10, 5))
  ax1 = fig.add_subplot(1, 2, 1)
  ax1.errorbar(freq, S, yerr=S_e,
            label ='Line1', color='red')
  ax1.set_xscale('log')
  ax1.set_yscale('log')
  ax1.set_ylabel('Spectral flux (mJy)')
  ax1.set_xticks([freq[0], freq[1]*0.4,freq[1]])
  #ax1.set_yticks([S[0], S[1]*0.4, S[1]])
  ax1.set_xlabel('frequency (MHz)')
  #slope = (S[1]-S[0])/(freq[1]-freq[0])
  anchored_text = AnchoredText(f'spectral index: {si[0]}', loc=1)
  ax1.add_artist(anchored_text)
  
  save_fig(plt,fig,kind,output=label)
  
def pl_RGB(ax, img,title,name,annot=True):
  """
    *Inputs*
    
        :rows:
          (int)
          Total number of rows.
        :columns:
          (int)
          Total number of columns.
        :i:
          (int)
          current number of the cell.
        :wcs:
          (astropy wcs)
          world coordinate system fetched from header of fits
        :img:
          (np array)
          dimension = (px,px,3)
        :fig:
          <maplotlib.pyplot.figure>
        :name:
          (string)
          input name to be show on figure

    *Returns*
    
        adds plot to the figure.
  """
  #--- RGBC plot -------------------#--#

  ax.axis('off')
  ax.imshow(img, origin='lower', cmap='gist_gray')
  if annot:
      ax.annotate("#RADatHomeIndia",(10,10),color='white')
      ax.annotate("By " + str(name),(400-5*len(name),10),color='white')
      ax.set_title('{}'.format(title),
                    y=1, pad=-16, color="white")
  ax.set_autoscale_on(False)

def overlayo(ri, gi, bi, kind = 'IOU'):
  """
  Returns RGB stacked image.

  *Input*

    :ri/gi/bi:
      (2-D array)
      survey 2-D array data in
    :kind:
      (string)
      either IOU or Optical

  *Returns*

    :Img:
      (nd array)
      dimension: (px,px,3)

    :scaling:
        - IOU
          - sqrt(w22), sqrt(dss2r), log(gnuv)(5 to 100% & factor =3.15)
        - Optical
          - sqrt, sqrt, sqrt (min to max)
  """
  if kind == 'IOU':
    ri = sqrt(ri, scale_min=np.percentile(np.unique(ri),1.), scale_max=np.percentile(np.unique(ri),100.))
    gi = sqrt(gi, scale_min=np.percentile(np.unique(gi),1.), scale_max=np.percentile(np.unique(gi),100.))
    #gi = log(gi, scale_min=np.percentile(np.unique(gi),0.), scale_max=np.percentile(np.unique(gi),100.),factor=2.85)
    bi = log(bi, scale_min=np.percentile(np.unique(bi), 5.),
             scale_max=np.percentile(np.unique(bi), 100.), factor=3.15)
    mul_factor = 255/ri.max()
    img = (np.transpose([(ri*mul_factor).astype(np.uint8),(gi*255/gi.max()).astype(np.uint8),(bi*255).astype(np.uint8)], (1, 2, 0)))

  if kind == 'Optical':
    #ri = log(ri, scale_min=np.percentile(np.unique(ri),1.), scale_max=np.percentile(np.unique(ri),100.),factor=0.3)
    ri = sqrt(ri, scale_min=np.min(ri), scale_max=np.percentile(np.unique(ri),100.))
    gi = sqrt(gi, scale_min=1.15*np.min(gi), scale_max=np.percentile(np.unique(gi),100.))
    #gi = log(gi, scale_min=np.percentile(np.unique(gi),0.), scale_max=np.percentile(np.unique(gi),100.),factor=7.85)
    #bi = log(bi, scale_min=np.percentile(np.unique(bi),1.), scale_max=np.percentile(np.unique(bi),100.),factor=3.15)
    bi = sqrt(bi, scale_min=np.min(bi), scale_max=np.percentile(np.unique(bi),100.))
    mul_factor = 255/ri.max()
    img = (np.transpose([(ri*mul_factor).astype(np.uint8),(gi*255.99).astype(np.uint8),(bi*256).astype(np.uint8)], (1, 2, 0)))
  #img = (np.dstack((ri,gi,bi))*255.99).astype(np.uint8)
  #img = np.stack([ri,gi,bi], axis=2)
  return img

def pl_RGBC(rows,columns,i,wcs,svy,lvlc,img,fig,name, pkind='ror',annot=True) :
    """
    *Inputs*

        :rows:
          (int)
          Total number of rows.
        :columns:
          (int)
          Total number of columns.
        :i:
          (int)
          current number of the cell.
        :wcs:
          (astropy wcs)
          world coordinate system fetched from header of fits
        :svy:
          (np array)
          2-D array survey for contours.
        :lvlc:
          (list)
          list of float
        :img:
          (np array)
          dimension = (px,px,3)
        :fig:
          <maplotlib.pyplot.figure>
        :name:
          (string)
          input name to be show on figure
        :pkind:
          (str)
          ror/iou;
          default='ror'

    *Returns*
    
        plots the ror/iou figure.
    """
    ax = fig.add_subplot(rows, columns, i, projection=wcs)
    ax.axis( 'off')
    ax.imshow(img, origin='lower', interpolation='nearest')
    if annot:
      ax.annotate("#RADatHomeIndia",(10,10),color='white')
      ax.annotate("By " + str(name),((400-5*len(name)),10),color='white')
    ax.set_autoscale_on(False)

    plt.contour(svy, lvlc, colors='white')
    if pkind == 'ror' and annot:
      if i==1 :
          ax.set_title("ROR-RGB-C: TGSS(GMRT)-DSS2-NVSS(VLA)-NVSS",
                      y=1, pad=-16, color="white")
          #ax.set_title("")
      if i==2 :
          ax.set_title("ROR-RGB-C: TGSS(GMRT)-DSS2-NVSS(VLA)-TGSS",
                      y=1, pad=-16, color="white")
    elif pkind == 'iou' and annot:
      if i==1 :
          ax.set_title("IOU-RGB-C: WISE(22)-DSS2(red)-GALEX(NUV)-TGSS",
                      y=1, pad=-16, color="white")
      if i==2 :
          ax.set_title("Optical-RGB-C: DSS2(IR)-DSS2(Red)-DSS2(blue)-TGSS",
                      y=1, pad=-16, color="white")


# ------------ arcsec to pixel conversion ----------------------------#-#

def to_pixel(unit_inarcsec, r, px = 480):
    """
Convert arcsec into pixel.
    
    :input:
      unit_inarcsec
        input magnitude (unitless) data in arcsec
      r
        length of square image in deg
      px
        total pixels in one length
    :Returns:
        input data converted to be used onto pixel length rounded upto 2 decimal places.
    """
    nvss_px_scale = px/(r)
    return np.round(((unit_inarcsec/3600) * ut.deg * nvss_px_scale), 2).value