import torch
import matplotlib.pyplot as plt
import random

class Synthetic_data:
    """
    This is a class for creating a lineal regression initialized with a given weight, bias and a noise
    generated by a normal distribution ~ N(0, 1)
    
    Attributes:
            w           (Tensor)   = the weight of linear regression
            b           (scalar)   = bias of the linear regression
            num_expales (int)      = lenght of the rows
            state       (string)   = can be either clean or noise -> clean reduce the standard deviation to 0.1         
    """
    def __init__(self, w, b, num_examples, state='noise'):
        self._w = w
        self._b = b
        self._num_examples = num_examples
        self._state = 'clean'
        self._features = None
        self._labels = None
        self._aleatoryMiniBatchFeatures = None
        self._aleatoryMiniBatchLabels = None
    
    def yield_data(self):
        
        """Return the features and labels of the multiple linear regression in two tensors, the state noise set a 
        epsilon cero, and noise set a epsilon of 0.1"""
        
        if self._state == 'noise':
            self._features = torch.normal(0, 1, (self._num_examples, len(self._w)))
            self._labels = torch.matmul(self._features,self._w) + self._b
            self._labels += torch.normal(0, 0.1, self._labels.shape)
            return self._features, self._labels.reshape((-1,1))
        
        elif self._state == 'clean':
            self._features = torch.normal(0, 1, (self._num_examples, len(self._w)))
            self._labels = torch.matmul(self._features,self._w) + self._b
            return self._features, self._labels.reshape((-1,1))         
        
    def data_iter(self,batch_size):  #(100, [1000,2], [1000,1])
        
        """shuffle the features and labels and return a generator of minibatches
        
            Parameters: 
                batch_size (int) = the lenght of each minibatch 
        """
        minibatch_features, minibatch_labels = self.yield_data()
        num_examples = len(minibatch_features)    # 1000
        indices = list(range(num_examples))  #[1000]
        random.shuffle(indices) 
        for i in range(0, num_examples, batch_size):  #->  0 1000, 100 -> 0, 100, ...900
            batch_indices = indices[i: min(i + batch_size, num_examples)]  # -> i: min(0+100, 1000)  ->  900 : min(900 + 100, 1000)
            yield self._features[batch_indices], self._labels[batch_indices] 
        
            
    def takeOneMinibatch(self,batch_size): 
        
        """Take one minibatch of the desired lenght and return it as two tensors (features and labels)
        
            Pararameter:
                batch_size (int) = lenght of the minibatch desired
                
        """
        count = 0
        flag = random.randint(0,int(self._num_examples//batch_size)-1)
        for X,y in self.data_iter(batch_size):
            if count == flag:
                return X,y
            count += 1