# В программе реализована простая линейная регрессия
# -*- coding: utf-8 -*-
"""
Created on 28 April, 2019
@author jswanglp

requirements:
    matplotlib==2.0.2
    numpy==1.15.4
    tensorflow==1.12.0
    scikit_learn==0.23.1

"""

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

if __name__ == '__main__':

    # Настройка параметров
    tf.app.flags.DEFINE_float('learning_rate', 3e-2, 'learning rate, default is 3e-2.')
    tf.app.flags.DEFINE_integer('num', 32, 'number of samples, default is 32.')
    tf.app.flags.DEFINE_integer('num_epochs', 50, 'number of epochs, default is 50.')
    FLAGS = tf.app.flags.FLAGS

    graph = tf.Graph()
    with graph.as_default():

        with tf.name_scope('Input'):
            x_input = tf.placeholder(tf.float32, shape=[None,], name='x_input')
            y_input = tf.placeholder(tf.float32, shape=[None,], name='y_input')
        with tf.name_scope('w_and_b'):
            w = tf.Variable(2.0, name='weight')
            b = tf.Variable(1.0, name='biases')
            y = tf.add(tf.multiply(x_input, w), b)
        with tf.name_scope('Train'):
            loss_op = tf.reduce_mean(tf.pow(y_input - y, 2))
            train_op = tf.train.AdamOptimizer(FLAGS.learning_rate).minimize(loss_op)

        gradients_node = tf.gradients(loss_op, w)

        # Настройка сессии
        sess = tf.Session()
        with sess.as_default():
            sess.run(tf.global_variables_initializer())

        # Генерация наборов обучающих данных
            x_pure = np.random.randint(-10, 100, FLAGS.num)
            x_train = x_pure + np.random.randn(FLAGS.num) / 10  # шум в направлении х
            y_train = 3 * x_pure + 2 + np.random.randn(FLAGS.num) / 10  # шум в направлении y
            Gradients = []
            Loss = []
            for i in range(FLAGS.num_epochs):
                _, gradients, loss = sess.run([train_op, gradients_node, loss_op],
                                              feed_dict={x_input: x_train,
                                                        y_input: y_train})
                print_list = [i+1, loss, gradients[0]]
                print("epoch: {0[0]} \t loss: {0[1]:.2f} \t gradients: {0[2]:.2f}".format(print_list))
                Gradients.append(gradients)
                Loss.append(loss)
        sess.close()

    fig = plt.figure(1, (14, 6))
    AX = [fig.add_subplot(i) for i in range(121,123)]
    name = ['Loss', 'Gradients']
    color = ['r', 'b']
    data = [Loss, Gradients]
    for na, ax, co, d in zip(name, AX, color, data):
        ax.plot(np.linspace(0, FLAGS.num_epochs, FLAGS.num_epochs), d, co, label=na)
        ax.set_title(na)
        ax.legend()

    plt.show()


"""Implementations of sequence alignment and optimal BST algorithms.

These algorithms provide examples of dynamic programming.

For now, I did not implement the reconstruction functions because the idea
is so similar to the algorithms in Chapter 16.
"""

import numpy as np

def read_seq_file(filename):
    """Reads data from sequence alignment test file.

    Args:
        filename (str): The file containing the edge list.

    Returns:
        str: The first sequence of characters.
        str: The second sequence of characters.
        int: The cost per gap in a sequence.
        int: The cost per mismatch in a sequence.
    """
    with open(filename, 'r') as f:
        next(f) # Skip first line
        cost_gap, cost_mismatch = next(f).strip().split()
        cost_gap, cost_mismatch = int(cost_gap), int(cost_mismatch)
        seq_x = next(f).strip()
        seq_y = next(f).strip()
        return seq_x, seq_y, cost_gap, cost_mismatch

def nw_score(seq_x, seq_y, cost_gap, cost_mismatch):
    """Returns the Needleman-Wunsch score for two sequences.

    The NW score is based on an optimal alignment, allowing gaps.

    I'm using a numpy array but avoiding vectorization to keep things minimal.

    Args:
        seq_x (str): The first sequence.
        seq_y (str): The second sequence.
        cost_gap (int): Cost per gap.
        cost_mismatch (int): Cost per mismatch.

    Returns:
        int: The NW score.
    """
    # Initialize array for holding solutions
    cache = np.zeros((len(seq_x)+1, len(seq_y)+1))
    # Base case #1 - seq_y is empty
    for i in range(len(seq_x) + 1):
        cache[i,0] = i * cost_gap
    # Base case #2 - seq_x is empty
    for j in range(len(seq_y) + 1):
        cache[0,j] = j * cost_gap
    # Systematically solve all sub-problems
    for i in range(1, len(seq_x) + 1):
        for j in range(1, len(seq_y) + 1):
            if seq_x[i-1] == seq_y[j-1]:
                match_cost = 0
            else:
                match_cost = cost_mismatch
            cost_1 = cache[i-1,j-1] + match_cost
            cost_2 = cache[i-1, j] + cost_gap
            cost_3 = cache[i, j-1] + cost_gap
            cache[i,j] = min(cost_1, cost_2, cost_3)
    return cache[len(seq_x), len(seq_y)]

def read_bst_file(filename):
    """Reads data from sequence alignment test file.

    Args:
        filename (str): The file containing the edge list.

    Returns:
        list: The frequencies for each key.
    """
    with open(filename, 'r') as f:
        next(f) # Skip first line
        freq = [int(x) for x in next(f).strip().split(',')]
        return freq

def opt_bst(freq):
    """Returns the search cost for an optimal BST.

    This algorithm finds an optimal BST, provided key frequencies.

    The costs are not normalized, so search cost is not average length.

    Args:
        freq (list): The frequencies for each key.

    Returns:
        int: The search cost for the optimal BST.
    """
    # Initialize array for holding solutions
    n = len(freq)
    cache = np.zeros((n+1, n+1))
    # Base case
    for i in range(n + 1):
        cache[i,i] = 0
    # Systematically solve all subproblems
    for s in range(n):
        for i in range(1, n + 1 - s):
            pk = sum(freq[i-1:i+s])
            min_cost = float('inf')
            for r in range(i, i + s + 1):
                min_cost = min(min_cost, cache[i-1,r-1] + cache[r,i+s])
            cache[i-1][i+s] = pk + min_cost
    return cache[0][n]

if __name__ == '__main__':
    import time

    nw_tests = [('problem17.8nw.txt', 224)]
    bst_opt_tests = [('problem17.8optbst.txt', 2780)]

    # Run NW correctness and timing tests
    print('Starting NW correctness tests...')
    for test_case in nw_tests:
        nw_inputs = read_seq_file(test_case[0])
        start = time.time()
        score = nw_score(*nw_inputs)
        end = round(time.time() - start, 3)
        assert score == test_case[1], f'Failed correctness for {test_case[0]}.'
        print(f'Passed correctness for {test_case[0]} in {end} seconds.')
    print('All NW correctness tests passed.')

    # Run NW correctness and timing tests
    print('Starting optimal BST correctness tests...')
    for test_case in bst_opt_tests:
        freq = read_bst_file(test_case[0])
        start = time.time()
        cost = opt_bst(freq)
        end = round(time.time() - start, 3)
        assert cost == test_case[1], f'Failed correctness for {test_case[0]}.'
        print(f'Passed correctness for {test_case[0]} in {end} seconds.')
    print('All optimal BST correctness tests passed.')


# Originally from https://cvw.cac.cornell.edu/python/exercise
# Modified by the cloudmesh team
from __future__ import print_function, division
"""
An estimate of the numerical value of pi via Monte Carlo integration.
Computation is distributed across processors via MPI.
"""

import numpy as np
from mpi4py import MPI
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import sys
from numba import jit
from cloudmesh.common.StopWatch import StopWatch

StopWatch.start("Overall time")
@jit(nopython=True)
def throw_darts(n):
    """
    returns an array of n uniformly random (x,y) pairs lying within the
    square that circumscribes the unit circle centered at the origin,
    i.e., the square with corners at (-1,-1), (-1,1), (1,1), (1,-1)
    """
    darts = 2*np.random.random((n,2)) - 1
    return darts

def in_unit_circle(p):
    """
    returns a boolean array, whose elements are True if the corresponding
    point in the array p is within the unit circle centered at the origin,
    and False otherwise -- hint: use np.linalg.norm to find the length of a vector
    """
    return np.linalg.norm(p,axis=-1)<=1.0

def estimate_pi(n, block=100000):
    """
    returns an estimate of pi by drawing n random numbers in the square
    [[-1,1], [-1,1]] and calculating what fraction land within the unit circle;
    in this version, draw random numbers in blocks of the specified size,
    and keep a running total of the number of points within the unit circle;
    by throwing darts in blocks, we are spared from having to allocate
    very large arrays (and perhaps running out of memory), but still can get
    good performance by processing large arrays of random numbers
    """
    total_number = 0
    i = 0
    while i < n:
        if n-i < block:
            block = n-i
        darts = throw_darts(block)
        number_in_circle = np.sum(in_unit_circle(darts))
        total_number += number_in_circle
        i += block
    return (4.*total_number)/n

def estimate_pi_in_parallel(comm, N):
    """
    on each of the available processes,
    calculate an estimate of pi by drawing N random numbers;
    the manager process will assemble all of the estimates
    produced by all workers, and compute the mean and
    standard deviation across the independent runs
    """

    if rank == 0:
        data = [N for i in range(size)]
    else:
        data = None
    data = comm.scatter(data, root=0)
    #
    pi_est = estimate_pi(N)
    #
    pi_estimates = comm.gather(pi_est, root=0)
    if rank == 0:
        return pi_estimates

def estimate_pi_statistics(comm, Ndarts, Nruns_per_worker):
    results = []
    for i in range(Nruns_per_worker):
        result = estimate_pi_in_parallel(comm, Ndarts)
        if rank == 0:
            results.append(result)
    if rank == 0:
        pi_est_mean = np.mean(results)
        pi_est_std  = np.std(results)
        return pi_est_mean, pi_est_std

if __name__ == '__main__':
    """
    for N from 4**5 to 4**14 (integer powers of 4),
    compute mean and standard deviation of estimates of pi
    by throwing N darts multiple times (Nruns_total times,
    distributed across workers)
    """
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()
    if rank == 0:
        print("MPI size = {}".format(size))
        sys.stdout.flush()
    Nruns_total = 64
    Nruns_per_worker = Nruns_total // size
    #
    estimates = []
    for log4N in range(5,15):
        N = int(4**log4N)
        result = estimate_pi_statistics(comm, N, Nruns_per_worker)
        if rank == 0:
            pi_est_mean, pi_est_std = result
            estimates.append((N, pi_est_mean, pi_est_std))
            print(N, pi_est_mean, pi_est_std)
            sys.stdout.flush()
    if rank == 0:
        estimates = np.array(estimates)
        plt.figure()
        plt.errorbar(np.log2(estimates[:,0]), estimates[:,1], yerr=estimates[:,2])
        plt.ylabel('estimate of pi')
        plt.xlabel('log2(number of darts N)')
        plt.savefig('pi_vs_log2_N.png')
        plt.figure()
        plt.ylabel('log2(standard deviation)')
        plt.xlabel('log2(number of darts N)')
        plt.plot(np.log2(estimates[:,0]), np.log2(estimates[:,2]))
        plt.savefig('log2_std_vs_log2_N.png')
    MPI.Finalize()

StopWatch.stop("Overall time")
StopWatch.benchmark()