# credit source 

import getopt
import sys

import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib import ticker
from matplotlib.colors import LinearSegmentedColormap

class NDVI(object):
    def __init__(self, file_path, output_file=False, colors=False):
        self.image = plt.imread(file_path)
        self.output_name = output_file or 'NDVI.jpg'
        self.colors = colors or ['red', 'orange', 'yellow', 'green', 'blue']

    def create_colormap(self, *args):
        return LinearSegmentedColormap.from_list(name='custom1', colors=args)

    def convert(self):
        """
        This function performs the NDVI calculation and returns an GrayScaled frame with mapped colors)
        """
        NIR = (self.image[:, :, 0]).astype('float')
        blue = (self.image[:, :, 2]).astype('float')
        green = (self.image[:, :, 1]).astype('float')
        bottom = (blue - green) ** 2
        bottom[bottom == 0] = 1  # replace 0 from nd.array with 1
        VIS = (blue + green) ** 2 / bottom
        NDVI = (NIR - VIS) / (NIR + VIS)
        plt.imsave(fname=self.output_name, arr =NDVI, cmap=self.create_colormap(*self.colors))



class VARI(object):
    def __init__(self, file_path, output_file, colors=False):
        self.image = plt.imread(file_path)
        self.output_file = output_file
        self.colors = colors or ['red', 'orange', 'yellow', 'green', 'blue']

    def create_colormap_VARI(self, *args):
        return LinearSegmentedColormap.from_list(name='custom1', colors=args)

    def convert_VARI(self):
        red = (self.image[:, :, 0]).astype('float')
        blue = (self.image[:, :, 2]).astype('float')
        green = (self.image[:, :, 1]).astype('float')
        VARI = (green - red) / (red + green - blue + .001)
        plt.imsave(fname=self.output_file, arr=VARI, cmap=self.create_colormap_VARI(*self.colors))


class TGI(object):
    def __init__(self, file_path, output_file=False, colors=False):
        self.image = plt.imread(file_path)
        self.output_file = output_file
        self.colors = colors or ['red', 'orange', 'yellow', 'green', 'blue']

    def create_colormap_TGI(self, *args):
        return LinearSegmentedColormap.from_list(name='custom1', colors=args)

    def convert_TGI(self):
        red = (self.image[:, :, 0]).astype('float')
        blue = (self.image[:, :, 2]).astype('float')
        green = (self.image[:, :, 1]).astype('float')
        TGI = (green - (.39 *red)) - (.61 * blue)
        plt.imsave(fname=self.output_file, arr=TGI, cmap=self.create_colormap_TGI(*self.colors))