import os
import random
import shutil
import time
import whisper

from django.conf import settings
from django.test import override_settings
from mock import patch, Mock

from .base import TestCase

from graphite.finders.utils import BaseFinder
from graphite.intervals import Interval, IntervalSet
from graphite.node import LeafNode, BranchNode
from graphite.readers.utils import BaseReader
from graphite.storage import Store, extractForwardHeaders, get_finders, get_tagdb, write_index
from graphite.tags.localdatabase import LocalDatabaseTagDB
from graphite.worker_pool.pool import PoolTimeoutError
from graphite.render.datalib import TimeSeries
from graphite.render.evaluator import evaluateTarget
from graphite.util import epoch_to_dt

class StorageTest(TestCase):

  def test_fetch(self):
    disabled_finder = get_finders('tests.test_storage.DisabledFinder')[0]
    legacy_finder = get_finders('tests.test_storage.LegacyFinder')[0]
    test_finder = get_finders('tests.test_storage.TestFinder')[0]
    remote_finder = get_finders('tests.test_storage.RemoteFinder')[0]

    store = Store(
      finders=[disabled_finder, legacy_finder, test_finder, remote_finder],
      tagdb=get_tagdb('graphite.tags.localdatabase.LocalDatabaseTagDB')
    )

    # tagb is properly initialized
    self.assertIsInstance(store.tagdb, LocalDatabaseTagDB)

    # get all enabled finders
    finders = store.get_finders()
    self.assertEqual(list(finders), [legacy_finder, test_finder, remote_finder])

    # get only local finders
    finders = store.get_finders(local=True)
    self.assertEqual(list(finders), [legacy_finder, test_finder])

    # fetch with empty patterns
    result = store.fetch([], 1, 2, 3, {})
    self.assertEqual(result, [])

    # fetch
    result = store.fetch(['a.**'], 1, 2, 3, {})
    self.assertEqual(len(result), 3)
    result.sort(key=lambda node: node['name'])
    self.assertEqual(result[0]['name'], 'a.b.c.d')
    self.assertEqual(result[0]['pathExpression'], 'a.**')
    self.assertEqual(result[1]['name'], 'a.b.c.d')
    self.assertEqual(result[1]['pathExpression'], 'a.**')
    self.assertEqual(result[2]['name'], 'a.b.c.e')
    self.assertEqual(result[2]['pathExpression'], 'a.**')

  def test_fetch_pool_timeout(self):
    # pool timeout
    store = Store(
      finders=[RemoteFinder()]
    )

    def mock_pool_exec(pool, jobs, timeout):
      raise PoolTimeoutError()

    message = r'Timed out after [-.e0-9]+s for fetch for \[\'a\'\]'
    with patch('graphite.storage.pool_exec', mock_pool_exec):
      with patch('graphite.storage.log.info') as log_info:
        with self.assertRaisesRegexp(Exception, message):
          list(store.fetch(['a'], 1, 2, 3, {}))
        self.assertEqual(log_info.call_count, 1)
        self.assertRegexpMatches(log_info.call_args[0][0], message)

  def test_fetch_all_failed(self):
    # all finds failed
    store = Store(
      finders=[TestFinder()]
    )

    message = r'All requests failed for fetch for \[\'a\'\] \(1\)'
    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, message):
        list(store.fetch(['a'], 1, 2, 3, {}))
      self.assertEqual(log_info.call_count, 1)
      self.assertRegexpMatches(
              log_info.call_args[0][0],
              r'Exception during fetch for \[\'a\'\] after [-.e0-9]+s: TestFinder.find_nodes'
      )

    store = Store(
      finders=[TestFinder(), TestFinder()]
    )

    message = r'All requests failed for fetch for \[\'a\'\] \(2\)'
    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, message):
        list(store.fetch(['a'], 1, 2, 3, {}))
      self.assertEqual(log_info.call_count, 2)
      self.assertRegexpMatches(
              log_info.call_args[0][0],
              r'Exception during fetch for \[\'a\'\] after [-.e0-9]+s: TestFinder.find_nodes'
      )

  def test_fetch_some_failed(self):
    # some finders failed
    store = Store(
      finders=[TestFinder(), RemoteFinder()]
    )

    with patch('graphite.storage.log.info') as log_info:
      list(store.fetch(['a'], 1, 2, 3, {}))
      self.assertEqual(log_info.call_count, 1)

    store = Store(
      finders=[TestFinder(), TestFinder()]
    )

    message = r'All requests failed for fetch for \[\'a\'\] \(2\)'
    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, message):
        list(store.fetch(['a'], 1, 2, 3, {}))
      self.assertEqual(log_info.call_count, 2)
      self.assertRegexpMatches(
              log_info.call_args[0][0],
              r'Exception during fetch for \[\'a\'\] after [-.e0-9]+s: TestFinder.find_nodes'
      )

  @override_settings(STORE_FAIL_ON_ERROR=True)
  def test_fetch_some_failed_hard_fail_enabled(self):
    # all finds failed
    store = Store(
      finders=[TestFinder(), RemoteFinder()]
    )

    message = r'1 request\(s\) failed for fetch for \[\'a\'\] \(2\)'
    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, message):
        list(store.fetch(['a'], 1, 2, 3, {}))
      self.assertEqual(log_info.call_count, 1)
      self.assertRegexpMatches(
              log_info.call_args[0][0],
              r'Exception during fetch for \[\'a\'\] after [-.e0-9]+s: TestFinder.find_nodes'
      )

    store = Store(
      finders=[TestFinder(), TestFinder()]
    )

    message = r'All requests failed for fetch for \[\'a\'\] \(2\)'
    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, message):
        list(store.fetch(['a'], 1, 2, 3, {}))
      self.assertEqual(log_info.call_count, 2)
      self.assertRegexpMatches(
              log_info.call_args[0][0],
              r'Exception during fetch for \[\'a\'\] after [-.e0-9]+s: TestFinder.find_nodes'
      )

  def test_find(self):
    disabled_finder = DisabledFinder()
    legacy_finder = LegacyFinder()
    test_finder = TestFinder()
    remote_finder = RemoteFinder()

    store = Store(
      finders=[disabled_finder, legacy_finder, test_finder, remote_finder],
      tagdb=get_tagdb('graphite.tags.localdatabase.LocalDatabaseTagDB')
    )

    # find nodes
    result = list(store.find('a'))
    self.assertEqual(len(result), 5)

    for node in result:
      if node.path in ['a.b.c.d', 'a.b.c.e']:
        self.assertIsInstance(node, LeafNode)
      else:
        self.assertIsInstance(node, BranchNode)
        self.assertTrue(node.path in ['a', 'a.b', 'a.b.c'])

    # find leaves only
    result = list(store.find('a', leaves_only=True))
    self.assertEqual(len(result), 2)

    for node in result:
      self.assertIsInstance(node, LeafNode)
      self.assertTrue(node.path in ['a.b.c.d', 'a.b.c.e'])

    # failure threshold
    message = r'Query a yields too many results and failed \(failure threshold is 1\)'
    with self.settings(METRICS_FIND_FAILURE_THRESHOLD=1):
      with self.assertRaisesRegexp(Exception, message):
        list(store.find('a'))

    # warning threshold
    with self.settings(METRICS_FIND_WARNING_THRESHOLD=1):
      with patch('graphite.storage.log.warning') as log_warning:
        list(store.find('a'))
        self.assertEqual(log_warning.call_count, 1)
        self.assertEqual(
          log_warning.call_args[0][0],
          'Query a yields large number of results up to 2 (warning threshold is 1)'
        )

  def test_find_pool_timeout(self):
    # pool timeout
    store = Store(
      finders=[RemoteFinder()]
    )

    def mock_pool_exec(pool, jobs, timeout):
      raise PoolTimeoutError()

    message = r'Timed out after [-.e0-9]+s for find <FindQuery: a from \* until \*>'
    with patch('graphite.storage.pool_exec', mock_pool_exec):
      with patch('graphite.storage.log.info') as log_info:
        with self.assertRaisesRegexp(Exception, message):
          list(store.find('a'))
        self.assertEqual(log_info.call_count, 1)
        self.assertRegexpMatches(log_info.call_args[0][0], message)

  def test_find_all_failed(self):
    # all finds failed
    store = Store(
      finders=[TestFinder()]
    )

    message = r'All requests failed for find <FindQuery: a from \* until \*>'
    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, message):
        list(store.find('a'))
      self.assertEqual(log_info.call_count, 1)
      self.assertRegexpMatches(
        log_info.call_args[0][0],
        r'Exception during find <FindQuery: a from \* until \*> after [-.e0-9]+s: TestFinder.find_nodes'
      )

    store = Store(
      finders=[TestFinder(), TestFinder()]
    )

    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, message):
        list(store.find('a'))
      self.assertEqual(log_info.call_count, 2)
      self.assertRegexpMatches(
        log_info.call_args[0][0],
        r'Exception during find <FindQuery: a from \* until \*> after [-.e0-9]+s: TestFinder.find_nodes'
      )

  @override_settings(REMOTE_STORE_FORWARD_HEADERS=['X-Test1', 'X-Test2'])
  def test_extractForwardHeaders(self):
    class DummyRequest(object):
      META = {
        'HTTP_X_TEST1': 'test',
      }

    headers = extractForwardHeaders(DummyRequest())
    self.assertEqual(headers, {'X-Test1': 'test'})

  def test_get_index(self):
    disabled_finder = DisabledFinder()
    # use get_finders so legacy_finder is patched with get_index
    legacy_finder = get_finders('tests.test_storage.LegacyFinder')[0]
    test_finder = TestFinder()
    remote_finder = RemoteFinder()

    store = Store(
      finders=[disabled_finder, legacy_finder, test_finder, remote_finder],
      tagdb=get_tagdb('graphite.tags.localdatabase.LocalDatabaseTagDB')
    )

    # get index
    result = store.get_index()
    self.assertEqual(result, ['a.b.c.d', 'a.b.c.e'])

    # get local index
    result = store.get_index({'localOnly': True})
    self.assertEqual(result, ['a.b.c.d'])

  def test_get_index_pool_timeout(self):
    # pool timeout
    store = Store(
      finders=[RemoteFinder()]
    )

    def mock_pool_exec(pool, jobs, timeout):
      raise PoolTimeoutError()

    with patch('graphite.storage.pool_exec', mock_pool_exec):
      with patch('graphite.storage.log.info') as log_info:
        with self.assertRaisesRegexp(Exception, 'Timed out after .*'):
          store.get_index()
        self.assertEqual(log_info.call_count, 1)
        self.assertRegexpMatches(log_info.call_args[0][0], 'Timed out after [-.e0-9]+s')

  def test_get_index_all_failed(self):
    # all finders failed
    store = Store(
      finders=[TestFinder()]
    )

    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, 'All requests failed for get_index'):
        store.get_index()
      self.assertEqual(log_info.call_count, 1)
      self.assertRegexpMatches(
              log_info.call_args[0][0],
              'Exception during get_index after [-.e0-9]+s: TestFinder.find_nodes'
      )

    store = Store(
      finders=[TestFinder(), TestFinder()]
    )

    with patch('graphite.storage.log.info') as log_info:
      with self.assertRaisesRegexp(Exception, r'All requests failed for get_index \(2\)'):
        store.get_index()
      self.assertEqual(log_info.call_count, 2)
      self.assertRegexpMatches(
              log_info.call_args[0][0],
              'Exception during get_index after [-.e0-9]+s: TestFinder.find_nodes'
      )

  @override_settings(USE_WORKER_POOL=False)
  def test_fetch_tag_support(self):
    class TestFinderTags(BaseFinder):
      tags = True

      def find_nodes(self, query):
        pass

      def fetch(self, patterns, start_time, end_time, now=None, requestContext=None):
        if patterns != [
                'seriesByTag("hello=tiger")',
                'seriesByTag("name=notags")',
                'seriesByTag("name=testtags")',
                'testtags;hello=tiger']:
          raise Exception('Unexpected patterns %s' % str(patterns))

        return [
          {
            'pathExpression': 'testtags;hello=tiger',
            'name': 'testtags;hello=tiger',
            'time_info': (0, 60, 1),
            'values': [],
          },
          {
            'pathExpression': 'seriesByTag("hello=tiger")',
            'name': 'testtags;hello=tiger',
            'time_info': (0, 60, 1),
            'values': [],
          },
          {
            'pathExpression': 'seriesByTag("name=testtags")',
            'name': 'testtags;hello=tiger',
            'time_info': (0, 60, 1),
            'values': [],
          },
        ]

    tagdb = Mock()

    store = Store(
      finders=[TestFinderTags()],
      tagdb=tagdb
    )

    request_context = {
      'startTime': epoch_to_dt(0),
      'endTime': epoch_to_dt(60),
      'now': epoch_to_dt(60),
    }

    with patch('graphite.render.datalib.STORE', store):
      results = evaluateTarget(request_context, ['testtags;hello=tiger', 'seriesByTag("hello=tiger")', 'seriesByTag("name=testtags")', 'seriesByTag("name=notags")'])
      self.assertEqual(results, [
        TimeSeries('testtags;hello=tiger', 0, 60, 1, []),
        TimeSeries('testtags;hello=tiger', 0, 60, 1, [], pathExpression='seriesByTag("hello=tiger")'),
        TimeSeries('testtags;hello=tiger', 0, 60, 1, [], pathExpression='seriesByTag("name=testtags")'),
      ])

  @override_settings(USE_WORKER_POOL=True)
  def test_fetch_no_tag_support(self):
    class TestFinderNoTags(BaseFinder):
      tags = False

      def find_nodes(self, query):
        pass

      def fetch(self, patterns, start_time, end_time, now=None, requestContext=None):
        if patterns != ['notags;hello=tiger']:
          raise Exception('Unexpected patterns %s' % str(patterns))

        return [
          {
            'pathExpression': 'notags;hello=tiger',
            'name': 'notags;hello=tiger',
            'time_info': (0, 60, 1),
            'values': [],
          }
        ]

    tagdb = Mock()

    def mockFindSeries(exprs, requestContext=None):
      self.assertEqual(requestContext, request_context)
      if exprs == ('hello=tiger',) or exprs == ('name=notags',):
        return ['notags;hello=tiger']
      if exprs == ('name=testtags',):
        return []
      raise Exception('Unexpected exprs %s' % str(exprs))

    tagdb.find_series.side_effect = mockFindSeries

    store = Store(
      finders=[TestFinderNoTags()],
      tagdb=tagdb
    )

    with patch('graphite.render.datalib.STORE', store):
      request_context = {
        'startTime': epoch_to_dt(0),
        'endTime': epoch_to_dt(60),
        'now': epoch_to_dt(60),
      }

      results = evaluateTarget(request_context, ['notags;hello=tiger', 'seriesByTag("hello=tiger")', 'seriesByTag("name=testtags")', 'seriesByTag("name=notags")'])
      self.assertEqual(tagdb.find_series.call_count, 3)
      self.assertEqual(results, [
        TimeSeries('notags;hello=tiger', 0, 60, 1, []),
        TimeSeries('notags;hello=tiger', 0, 60, 1, [], pathExpression='seriesByTag("hello=tiger")'),
        TimeSeries('notags;hello=tiger', 0, 60, 1, [], pathExpression='seriesByTag("name=notags")'),
      ])

  def test_autocomplete(self):
    test = self

    class TestFinderTags(BaseFinder):
      tags = True

      def __init__(self, request_limit=100, request_context=None):
        self.limit = request_limit
        self.context = request_context or {}

      def find_nodes(self, query):
        pass

      def auto_complete_tags(self, exprs, tagPrefix=None, limit=None, requestContext=None):
        test.assertEqual(exprs, ['tag1=value1'])
        test.assertEqual(tagPrefix, 'test')
        test.assertEqual(limit, self.limit)
        test.assertEqual(requestContext, self.context)
        return ['testtags']

      def auto_complete_values(self, exprs, tag, valuePrefix=None, limit=None, requestContext=None):
        test.assertEqual(exprs, ['tag1=value1'])
        test.assertEqual(tag, 'tag2')
        test.assertEqual(valuePrefix, 'test')
        test.assertEqual(limit, self.limit)
        test.assertEqual(requestContext, self.context)
        return ['testtags']

    class TestFinderNoTags(BaseFinder):
      tags = False

      def find_nodes(self, query):
        pass

    class TestFinderTagsException(TestFinderTags):
      def auto_complete_tags(self, exprs, tagPrefix=None, limit=None, requestContext=None):
        raise Exception('TestFinderTagsException.auto_complete_tags')

      def auto_complete_values(self, exprs, tag, valuePrefix=None, limit=None, requestContext=None):
        raise Exception('TestFinderTagsException.auto_complete_values')

    class TestFinderTagsTimeout(TestFinderTags):
      def auto_complete_tags(self, exprs, tagPrefix=None, limit=None, requestContext=None):
        time.sleep(0.1)
        return ['testtags']

      def auto_complete_values(self, exprs, tag, valuePrefix=None, limit=None, requestContext=None):
        time.sleep(0.1)
        return ['testtags']

    def mockStore(finders, request_limit=100, request_context=None):
      tagdb = Mock()

      def mockAutoCompleteTags(exprs, tagPrefix=None, limit=None, requestContext=None):
        self.assertEqual(exprs, ['tag1=value1'])
        self.assertEqual(tagPrefix, 'test')
        self.assertEqual(limit, request_limit)
        self.assertEqual(requestContext, request_context or {})
        return ['testnotags']

      tagdb.auto_complete_tags.side_effect = mockAutoCompleteTags

      def mockAutoCompleteValues(exprs, tag, valuePrefix=None, limit=None, requestContext=None):
        self.assertEqual(exprs, ['tag1=value1'])
        self.assertEqual(tag, 'tag2')
        self.assertEqual(valuePrefix, 'test')
        self.assertEqual(limit, request_limit)
        self.assertEqual(requestContext, request_context or {})
        return ['testnotags']

      tagdb.auto_complete_values.side_effect = mockAutoCompleteValues

      return Store(
        finders=finders,
        tagdb=tagdb,
      )

    request_context = {}

    # test with both tag-enabled and non-tag-enabled finders
    store = mockStore([TestFinderTags(), TestFinderNoTags()])

    result = store.tagdb_auto_complete_tags(['tag1=value1'], 'test', 100, request_context)
    self.assertEqual(store.tagdb.auto_complete_tags.call_count, 1)
    self.assertEqual(result, ['testnotags', 'testtags'])

    result = store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test', 100, request_context)
    self.assertEqual(store.tagdb.auto_complete_values.call_count, 1)
    self.assertEqual(result, ['testnotags', 'testtags'])

    # test with no limit & no requestContext
    store = mockStore([TestFinderTags(None, {}), TestFinderNoTags()], None, {})

    result = store.tagdb_auto_complete_tags(['tag1=value1'], 'test')
    self.assertEqual(store.tagdb.auto_complete_tags.call_count, 1)
    self.assertEqual(result, ['testnotags', 'testtags'])

    result = store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test')
    self.assertEqual(store.tagdb.auto_complete_values.call_count, 1)
    self.assertEqual(result, ['testnotags', 'testtags'])

    # test with only tag-enabled finder
    store = mockStore([TestFinderTags()])

    result = store.tagdb_auto_complete_tags(['tag1=value1'], 'test', 100, request_context)
    self.assertEqual(store.tagdb.auto_complete_tags.call_count, 0)
    self.assertEqual(result, ['testtags'])

    result = store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test', 100, request_context)
    self.assertEqual(store.tagdb.auto_complete_values.call_count, 0)
    self.assertEqual(result, ['testtags'])

    # test with only non-tag-enabled finder
    store = mockStore([TestFinderNoTags()])

    result = store.tagdb_auto_complete_tags(['tag1=value1'], 'test', 100, request_context)
    self.assertEqual(store.tagdb.auto_complete_tags.call_count, 1)
    self.assertEqual(result, ['testnotags'])

    result = store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test', 100, request_context)
    self.assertEqual(store.tagdb.auto_complete_values.call_count, 1)
    self.assertEqual(result, ['testnotags'])

    # test with no finders
    store = mockStore([])

    result = store.tagdb_auto_complete_tags(['tag1=value1'], 'test', 100, request_context)
    self.assertEqual(store.tagdb.auto_complete_tags.call_count, 0)
    self.assertEqual(result, [])

    result = store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test', 100, request_context)
    self.assertEqual(store.tagdb.auto_complete_values.call_count, 0)
    self.assertEqual(result, [])

    # test exception handling with one finder
    store = mockStore([TestFinderTagsException()])

    with self.assertRaisesRegexp(Exception, r'All requests failed for tags for \[\'tag1=value1\'\] test.*'):
      store.tagdb_auto_complete_tags(['tag1=value1'], 'test', 100, request_context)

    with self.assertRaisesRegexp(Exception, r'All requests failed for values for \[\'tag1=value1\'\] tag2 test.*'):
      store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test', 100, request_context)

    # test exception handling with more than one finder
    store = mockStore([TestFinderTagsException(), TestFinderTagsException()])

    with self.assertRaisesRegexp(Exception, r'All requests failed for tags for \[\'tag1=value1\'\] test'):
      store.tagdb_auto_complete_tags(['tag1=value1'], 'test', 100, request_context)

    with self.assertRaisesRegexp(Exception, r'All requests failed for values for \[\'tag1=value1\'\] tag2 test'):
      store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test', 100, request_context)

    # test pool timeout handling
    store = mockStore([TestFinderTagsTimeout()])

    with self.settings(USE_WORKER_POOL=True, FIND_TIMEOUT=0):
      with self.assertRaisesRegexp(Exception, r'Timed out after [-.e0-9]+s for tags for \[\'tag1=value1\'\]'):
        store.tagdb_auto_complete_tags(['tag1=value1'], 'test', 100, request_context)

      with self.assertRaisesRegexp(Exception, r'Timed out after [-.e0-9]+s for values for \[\'tag1=value1\'\] tag2 test'):
        store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test', 100, request_context)

  # test write_index
  hostcpu = os.path.join(settings.WHISPER_DIR, 'hosts/hostname/cpu.wsp')

  def create_whisper_hosts(self):
    worker1 = self.hostcpu.replace('hostname', 'worker1')
    worker2 = self.hostcpu.replace('hostname', 'worker2')
    bogus_file = os.path.join(settings.WHISPER_DIR, 'a/b/c/bogus_file.txt')

    try:
      os.makedirs(worker1.replace('cpu.wsp', ''))
      os.makedirs(worker2.replace('cpu.wsp', ''))
      os.makedirs(bogus_file.replace('bogus_file.txt', ''))
    except OSError:
      pass

    open(bogus_file, 'a').close()
    whisper.create(worker1, [(1, 60)])
    whisper.create(worker2, [(1, 60)])

    ts = int(time.time())
    whisper.update(worker1, 1, ts)
    whisper.update(worker2, 2, ts)

  def wipe_whisper_hosts(self):
    try:
      os.remove(self.hostcpu.replace('hostname', 'worker1'))
      os.remove(self.hostcpu.replace('hostname', 'worker2'))
      os.remove(os.path.join(settings.WHISPER_DIR, 'a/b/c/bogus_file.txt'))
      shutil.rmtree(self.hostcpu.replace('hostname/cpu.wsp', ''))
      shutil.rmtree(os.path.join(settings.WHISPER_DIR, 'a'))
    except OSError:
      pass

  def test_write_index(self):
    self.create_whisper_hosts()
    self.addCleanup(self.wipe_whisper_hosts)

    self.assertEqual(None, write_index() )
    self.assertEqual(None, write_index(settings.INDEX_FILE) )

class DisabledFinder(object):
  disabled = True

  def find_nodes(self, query):
    pass

class LegacyFinder(object):
  def find_nodes(self, query):
    yield BranchNode('a')
    yield BranchNode('a.b')
    yield BranchNode('a.b.c')
    yield LeafNode('a.b.c.d', DummyReader('a.b.c.d'))

class DummyReader(BaseReader):
    __slots__ = ('path',)

    def __init__(self, path):
        self.path = path

    def fetch(self, startTime, endTime, now=None, requestContext=None):
        npoints = (endTime - startTime) // 10
        return (startTime, endTime, 10), [
            random.choice([None, 1, 2, 3]) for i in range(npoints)
        ]

    def get_intervals(self):
        return IntervalSet([Interval(time.time() - 3600, time.time())])

class RemoteFinder(BaseFinder):
  local = False

  def find_nodes(self, query):
    yield BranchNode('a.b.c')
    yield LeafNode('a.b.c.d', DummyReader('a.b.c.d'))
    yield LeafNode('a.b.c.e', DummyReader('a.b.c.e'))

class TestFinder(BaseFinder):
  def find_nodes(self, query):
    raise Exception('TestFinder.find_nodes')


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Feb 23 15:31:46 2018
script python to plot some outputs of the model
@author: jbrlod
modified by mkouassi & arimoux
"""
import os
import numpy as np
import matplotlib.pyplot as plt
import shutil
import xarray as xr
from baseutil import dataset, weights_mask

alltestname = ['BaseTest_Cloud1.nc','BaseTest_Clouds.nc','BaseTest_Square1.nc','BaseTest_Squares.nc']
#alltestname = ['base_validation_1cloud.nc', 'base_validation_multiple_clouds.nc', 'base_validation_1square.nc','base_validation_multiple_squares.nc']
## DATA DIRECTORY
AllInputDs = ["cl1", "cl+", "sq1", "sq+"]              # cl1 : cloud ; sq1: nauge carré ; sq2 : nuage carré centré
allPattern = ['cloud1','clouds','square1','squares'];  # list of the cloud's patterns
alldir = ['cloud1','clouds','square1','squares'];      # list of all data directories
AllDeepModels = ['cnn','gan']                          # lists of deep learning algorithms
deepmodels=AllDeepModels[0]
count = 0;
for icode, inputds in enumerate(AllInputDs):
    if (inputds==AllInputDs[0]):
        ddir = alldir[0]
        datadir = '../data/cloud/'+ddir
        pattern = allPattern[0]
        testname = alltestname[0]
        weightsName = 'Weights_Cloud1.nc'

    elif (inputds == AllInputDs[1]):
        ddir = alldir[1]
        datadir = '../data/cloud/'+ddir
        pattern = allPattern[1]
        testname = alltestname[1]
        weightsName = 'Weights_Clouds.nc'

    elif (inputds == AllInputDs[2]):
        ddir = alldir[2]
        datadir = '../data/square/'+ddir
        pattern = allPattern[2]
        testname = alltestname[2]
        weightsName = 'Weights_Square1.nc'

    elif (inputds == AllInputDs[3]):
        ddir = alldir[3]
        datadir = '../data/square/'+ddir
        pattern = allPattern[3]
        testname = alltestname[3]
        weightsName = 'Weights_Squares.nc'

    # CALL THE DATASET CLASS
    ds = dataset(basename=os.path.join(datadir,testname))

    # BUILDING of Binary input mask  and loss function weight masks
    Inputds = os.path.join(datadir,testname)
    weightdsName = os.path.join(datadir,weightsName)
    wds = weights_mask(Inputds, weightdsName, weight_c=0.1, weight_n=1)
    weight = wds.weights.expand_dims('canal',3)
    Xbinary = wds.bmask.expand_dims('canal',3)
    ## Concatenation des bases des dataArray 'yt',weights
    yt_2dim = xr.concat((ds.yt, weight),dim='canal')

    ## MODEL DIRECTORY
    from keras.models import load_model
    from modelutil import context_mse, masked_mse

    allLoss = ['mmse','cmse'];
    allLayers=['2layers','3layers','4layers']; nbLayers = allLayers[2]
    # list of models directories based on the inputs data models
    allmdir=['classic','binary']#, 'temporal' ,'bintemp']
    for mcode, mdir in enumerate(allmdir):
        modelsdir = '../data/models/'+deepmodels+'/'+mdir
        #
        if (mdir == allmdir[0]):
            Xbinary.values = np.array(Xbinary.values, dtype='int')
        elif (mdir == allmdir[1]):
            Xbinary.values = np.ones_like(Xbinary.values)

        X_2dim = xr.concat((ds.X, Xbinary),dim='canal')
        ## Splitting
        #X_test, X_valid, yt_test, yt_valid = train_test_split(X_2dim, yt_2dim, test_size=0.2)

        # Choice of the loss function
        for mcode, Lossf in enumerate(allLoss):
            # Name of the neural architecture
            name = 'model_'+nbLayers+'_'+pattern+'_'+Lossf+'_'+mdir
            # Outputname
            outname = 'DatasetNN_'+nbLayers+'_'+pattern+'_'+Lossf+'_'+mdir+'_'+deepmodels+'.nc'
            # Model loading
            if Lossf == allLoss[0]:
                model = load_model(os.path.join(modelsdir,name), custom_objects={'masked_mse':masked_mse})
            elif Lossf == allLoss[1]:
                model = load_model(os.path.join(modelsdir,name),custom_objects={'context_mse':context_mse})
            else:
                raise ValueError("Choisir une fonction de cout existante ou en créer une dans modelutil")

            ypredict = xr.DataArray(model.predict(X_2dim,batch_size=10,verbose=1),coords=ds.yt.coords)

            # Generate a combinated image of original & predicted images
            isCloud = np.equal(ds.X,0)
            yfinal = (1-isCloud)*ds.X + isCloud*ypredict

            # Save prediction
            dsout = xr.Dataset({'X':ds.X,'yt':ds.yt,'ypredict':ypredict,'yfinal':yfinal,'ytfull':ds.ytfull})
            dsout.to_netcdf(os.path.join(datadir,outname))

            print(count); count +=1

            # Plot some random images
            PLOT = False

            #save the images
            SAVE = False

            if PLOT:

                nim = 20 #number of images to be plot
                ii = np.random.randint(0,dsout.index.size,nim)
                vmin, vmax = -1.5,0

                for i,ind in enumerate(ii):
                    fig, axes= plt.subplots(ncols=4)
                    axes[0].imshow(np.log10(dsout.X[ind].squeeze()),vmin=vmin,vmax=vmax)
                    axes[1].imshow(np.log10(dsout.yt[ind].squeeze()),vmin=vmin,vmax=vmax)
                    axes[2].imshow(np.log10(dsout.ypredict[ind].squeeze()),vmin=vmin,vmax=vmax)
                    axes[3].imshow(np.log10(dsout.yfinal[ind].squeeze()),vmin=vmin,vmax=vmax)
                    title = 'Image_' + str(int(dsout.index[ind]))
                    plt.suptitle(title)
                    if SAVE:
                        # example dir : A FAIRE +++
                        exampledir = os.path.join('../figures/examples',os.path.splitext(name)[0])
                        if os.path.isdir():
                            print('Directory already created')
                            plt.savefig(os.path.join(exampledir,title+'.png'))
                        else:
                            shutil.rmtree(exampledir,ignore_errors=True)
                            os.makedirs(exampledir)
                            plt.savefig(os.path.join(exampledir,title+'.png'))


import mysql.connector
def my_conn() :
    config = {
      'user': 'root',
      'password': '^Kimjml5v2!',
      'host': '127.0.0.1',
      'database': 'my_emp',
      'raise_on_warnings': True
    }

    cnx = mysql.connector.connect(**config)
    print(cnx)
    return cnx

def select_All(conn):
    cur = conn.cursor()
    cur.execute("SELECT * FROM emp")    #return 없이 실행되는 구문이다.
    rows = cur.fetchall()
    for row in rows:
        print(row)

if __name__ == '__main__':
    select_All(my_conn())


import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
import warnings
from typing import Sequence

from numba import jit, prange

def leaf_samples(rf, X_not_col:np.ndarray) -> Sequence:
    """
    Return a list of arrays where each array is the set of X sample indexes
    residing in a single leaf of some tree in rf forest. For example, if there
    are 4 leaves (in one or multiple trees), we might return:

        array([array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
           array([10, 11, 12, 13, 14, 15]), array([16, 17, 18, 19, 20]),
           array([21, 22, 23, 24, 25, 26, 27, 28, 29]))
    """
    n_trees = len(rf.estimators_)
    leaf_samples = []
    leaf_ids = rf.apply(X_not_col)  # which leaf does each X_i go to for sole tree?
    for t in range(n_trees):
        # Group by id and return sample indexes
        uniq_ids = np.unique(leaf_ids[:,t])
        sample_idxs_in_leaves = [np.where(leaf_ids[:, t] == id)[0] for id in uniq_ids]
        leaf_samples.extend(sample_idxs_in_leaves)
    return leaf_samples

def partial_dependence(X:pd.DataFrame, y:pd.Series, colname:str,
                       min_slopes_per_x=5,
                       parallel_jit=True,
                       n_trees=1, min_samples_leaf=15, rf_bootstrap=False, max_features=1.0,
                       supervised=True,
                       verbose=False):
    """
    Internal computation of partial dependence information about X[colname]'s effect on y.
    Also computes partial derivative of y with respect to X[colname].

    :param X: Dataframe with all explanatory variables
    :param y: Series or vector with response variable
    :param colname: which X[colname] (a string) to compute partial dependence for
    :param min_slopes_per_x: ignore pdp y values derived from too few slopes; this is
           same count across all features (tried percentage of max slope count but was
           too variable). Important for getting good starting point of PD.

    Returns:
        leaf_xranges    The ranges of X[colname] partitions

        leaf_slopes     Associated slope for each leaf xrange

        slope_counts_at_x How many slopes are available at each x_j location

        dx              The change in x from one non-NaN unique X[colname] to the next

        slope_at_x      The slope at each non-NaN unique X[colname]

        pdpx            The non-NaN unique X[colname] values; len(pdpx)<=len(unique(X[colname]))

        pdpy            The effect of each non-NaN unique X[colname] on y; effectively
                        the cumulative sum (integration from X[colname] x to z for all
                        z in X[colname]). The first value is always 0.

        ignored         How many samples from len(X) total records did we have to
                        ignore because samples in leaves had identical X[colname]
                        values.
  """
    X_not_col = X.drop(colname, axis=1).values
    # For x floating-point numbers that are very close, I noticed that np.unique(x)
    # was treating floating-point numbers different in the 12th decimal point as different.
    # This caused a number of problems likely but I didn't notice it until I tried
    # np.gradient(), which found extremely huge derivatives. I fixed that with a hack:
    X_col = X[colname].values.round(decimals=10)

    if supervised:
        rf = RandomForestRegressor(n_estimators=n_trees,
                                   min_samples_leaf=min_samples_leaf,
                                   bootstrap=rf_bootstrap,
                                   max_features=max_features)
        rf.fit(X_not_col, y)
        if verbose:
            print(f"Strat Partition RF: dropping {colname} training R^2 {rf.score(X_not_col, y):.2f}")

    else:
        """
        Wow. Breiman's trick works in most cases. Falls apart on Boston housing MEDV target vs AGE
        """
        if verbose: print("USING UNSUPERVISED MODE")
        X_synth, y_synth = conjure_twoclass(X)
        rf = RandomForestClassifier(n_estimators=n_trees,
                                    min_samples_leaf=min_samples_leaf,
                                    bootstrap=rf_bootstrap,
                                    max_features=max_features,
                                    oob_score=False)
        rf.fit(X_synth.drop(colname, axis=1), y_synth)

    if verbose:
        leaves = leaf_samples(rf, X_not_col)
        nnodes = rf.estimators_[0].tree_.node_count
        print(f"Partitioning 'x not {colname}': {nnodes} nodes in (first) tree, "
              f"{len(rf.estimators_)} trees, {len(leaves)} total leaves")

    leaf_xranges, leaf_slopes, ignored = \
        collect_discrete_slopes(rf, X_col, X_not_col, y) # if ignored, won't have entries in leaf_* results

    # print('leaf_xranges', leaf_xranges)
    # print('leaf_slopes', leaf_slopes)

    real_uniq_x = np.unique(X_col)   # comes back sorted
    if verbose:
        print(f"discrete StratPD num samples ignored {ignored}/{len(X)} for {colname}")

    #print("uniq x =", len(real_uniq_x), "slopes.shape =", leaf_slopes.shape, "x ranges.shape", leaf_xranges.shape)
    if parallel_jit:
        slope_at_x, slope_counts_at_x = \
            avg_slopes_at_x_jit(real_uniq_x, leaf_xranges, leaf_slopes)
    else:
        slope_at_x, slope_counts_at_x = \
            avg_slopes_at_x_nonparallel_jit(real_uniq_x, leaf_xranges, leaf_slopes)

    if min_slopes_per_x <= 0:
        min_slopes_per_x = 1 # must have at least one slope value

    # Turn any slopes with weak evidence into NaNs but keep same slope_at_x length
    slope_at_x = np.where(slope_counts_at_x >= min_slopes_per_x, slope_at_x, np.nan)

    pdpx = real_uniq_x

    # At this point, slope_at_x will have at least one nan, but that's okay because
    # diff and nancumsum will treat them appropriately.  nancumsum treats NaN as 0
    # so the previous non-NaN value is carried forward until we reach next real value.
    # Keep in mind that the last slope in slope_at_x is always nan,
    # since we have no information beyond that. However, we can still produce
    # a pdpy value for that last position because we have a delta from the
    # previous position that will get us to the last position.

    # Integrate the partial derivative estimate in slope_at_x across pdpx to get dependence
    dx = np.diff(pdpx)      # for n pdpx values, there are n-1 dx values
    dydx = slope_at_x[:-1]  # for n y values, only n-1 slope values; last slope always nan

    y_deltas = dydx * dx                          # get change in y from x[i] to x[i+1]
    pdpy = np.nancumsum(y_deltas)                 # y_deltas has one less value than num x values
    pdpy = np.concatenate([np.array([0]), pdpy])  # align with x values; our PDP y always starts from zero

    # At this point pdpx, pdpy have the same length as real_uniq_x

    # Strip from pdpx,pdpy any positions for which we don't have useful slope info. If we have
    # slopes = [1, 3, nan] then cumsum will give 3 valid pdpy values. But, if we have
    # slopes = [1, 3, nan, nan], then there is no pdpy value for last position. nancumsum
    # carries same y to the right, but we wanna strip last position, in this case.
    # Detect 2 nans in a row. We don't want the second or any after that until next real value
    idx_not_adjacent_nans = np.where(~(np.isnan(slope_at_x[1:]) & np.isnan(slope_at_x[:-1])))[0] + 1
    # Deal with first position. If it's nan, drop it, else keep it
    if not np.isnan(slope_at_x[0]):
        idx_not_adjacent_nans = np.concatenate([np.array([0]), idx_not_adjacent_nans])
    pdpx = pdpx[idx_not_adjacent_nans]
    pdpy = pdpy[idx_not_adjacent_nans]

    return leaf_xranges, leaf_slopes, slope_counts_at_x, dx, slope_at_x, pdpx, pdpy, ignored

def finite_differences(x: np.ndarray, y: np.ndarray):
    """
    Use the unique x values within a leaf to compute finite differences. Given, n unique
    x values return n-1 derivative estimates. We start by grouping the leaf x,y by x
    and then collect the average y. The unique x and y averages are the new x and y pairs.

    The slope for each x is the forward discrete difference:

        (y_{i+1} - y_i) / (x_{i+1} - x_i)

    At right edge, there is no following x value so we could use backward difference:

        (y_i - y_{i-1}) / (x_{i+1} - x_i)

    But we don't use that last value as there's nothing to plot pass last value; slope
    at last value therefore is worthless.

    I experimented with the center difference using a secant method:

        (y_{i+1} - y_{i-1}) / (x_{i+1} - x_{i-1})

    which is Method A from Veldman, A.E.P., Rinzema, K. Playing with nonuniform grids.
    J Eng Math 26, 119–130 (1992).  https://doi.org/10.1007/BF00043231
    https://www.rug.nl/research/portal/files/3332271/1992JEngMathVeldman.pdf

    The np.gradient() function uses Method B from Veldman and Rinzema but that one seems
    to give much worse derivatives for sparse data.  In the end, forward difference seemed
    to get the best approximation for synthetic data, particularly sparse data. Since
    it is simplest, I went with it.

    If the ordinal/ints are exactly one unit part, then it's just y_{i+1} - y_i. If
    they are not consecutive, we do not ignore isolated x_i as it ignores too much data.
    E.g., if x is [1,3,4] and y is [9,8,10] then the x=2 coordinate is spanned as part
    of 1 to 3. The two slopes are [(8-9)/(3-1), (10-8)/(4-3)] and bin widths are [2,1].

    If there is exactly one unique x value in the leaf, the leaf provides no information
    about how X[colname] contributes to changes in y. We have to ignore this leaf.
    """
    ignored = 0

    # Group by x, take mean of all y with same x value (they come back sorted too)
    uniq_x = np.unique(x)
    avg_y = np.array([y[x==ux].mean() for ux in uniq_x])

    if len(uniq_x)==1:
        # print(f"ignore {len(x)} in discrete_xc_space")
        ignored += len(x)
        return np.array([[0]],dtype=x.dtype), np.array([0.0]), ignored

    # FORWARD DIFF
    x_deltas = np.diff(uniq_x)
    y_deltas = np.diff(avg_y)
    leaf_slopes = y_deltas / x_deltas  # "rise over run"

    # AVERAGE AROUND CENTER DIFF
    # At position i, take average of forward slope from y[i-1] to y[i] and
    # slope from y[i] to y[i+1].
    # leaf_slopes = [leaf_slopes[0]] + list((leaf_slopes[:-1] + leaf_slopes[1:]) / 2)
    # leaf_slopes = np.array(leaf_slopes)
    # dang, doesn't seem to work

    # CENTER DIFF
    # x_deltas2 = uniq_x[2:] - uniq_x[:-2]  # this is empty if |uniq_x|==2
    # y_deltas2 = avg_y[2:] - avg_y[:-2]
    # dydx0 = (avg_y[1] - avg_y[0]) / (uniq_x[1] - uniq_x[0]) # forward diff for i=0
    # leaf_slopes_ctr = [dydx0] + list(y_deltas2 / x_deltas2)  # "rise over run, stride 2"
    # leaf_slopes_ctr = np.array(leaf_slopes_ctr)

    # leaf_slopes = np.gradient(avg_y, uniq_x)[:-1] # drop last derivative as we won't use it
    # print("ctr",list(leaf_slopes_ctr))
    # print("grd",list(leaf_slopes))

    leaf_xranges = np.array(list(zip(uniq_x, uniq_x[1:])))

    return leaf_xranges, leaf_slopes, ignored

def collect_discrete_slopes(rf, X_col, X_not_col, y):
    """
    For each leaf of each tree of the decision tree or RF rf (trained on all features
    except colname), get the leaf samples then isolate the X[colname] values
    and the target y values.  Compute the y deltas between unique X[colname] values.
    Like performing piecewise linear regression of X[colname] vs y
    to get the slopes in various regions of X[colname].  We don't need to subtract
    the minimum y value before regressing because the slope won't be different.
    (We are ignoring the intercept of the regression line).

    Return for each leaf, the ranges of X[colname] partitions,
    associated slope for each range, and number of ignored samples.
    """
    # start = timer()
    leaf_slopes = []   # drop or rise between discrete x values
    leaf_xranges = []  # drop is from one discrete value to next

    ignored = 0

    leaves = leaf_samples(rf, X_not_col)
    y = y.values

    if False:
        nnodes = rf.estimators_[0].tree_.node_count
        print(f"Partitioning 'x not {colname}': {nnodes} nodes in (first) tree, "
              f"{len(rf.estimators_)} trees, {len(leaves)} total leaves")

    for samples in leaves:
        leaf_x = X_col[samples]
        # leaf_x = one_leaf_samples[]#.reshape(-1,1)
        leaf_y = y[samples]

        if np.abs(np.min(leaf_x) - np.max(leaf_x)) < 1.e-8: # faster than np.isclose()
            # print(f"ignoring xleft=xright @ {r[0]}")
            ignored += len(leaf_x)
            continue

        leaf_xranges_, leaf_slopes_, ignored_ = \
            finite_differences(leaf_x, leaf_y)

        leaf_slopes.extend(leaf_slopes_)
        leaf_xranges.extend(leaf_xranges_)
        ignored += ignored_

    if len(leaf_xranges)==0:
        # make sure empty list has same shape (jit complains)
        leaf_xranges = np.array([]).reshape(0, 0)
    else:
        leaf_xranges = np.array(leaf_xranges)
    leaf_slopes = np.array(leaf_slopes)

    # stop = timer()
    # if verbose: print(f"collect_discrete_slopes {stop - start:.3f}s")
    return leaf_xranges, leaf_slopes, ignored

# We get about 20% boost from parallel but limits use of other parallelism it seems;
# i get crashes when using multiprocessing package on top of this.
# If using n_jobs=1 all the time for importances, then turn jit=False so this
# method is not used
@jit(nopython=True, parallel=True) # use prange not range.
def avg_slopes_at_x_jit(uniq_x, leaf_ranges, leaf_slopes):
    """
    Compute the average of leaf_slopes at each uniq_x.

    Value at max(x) is NaN since we have no data beyond that point and so there is
    no forward difference. If last range is 4..5 then slope at 5 is nan since we
    don't know where it's going to go from there.
    """
    nx = uniq_x.shape[0]
    nslopes = leaf_slopes.shape[0]
    slopes = np.empty(shape=(nx, nslopes), dtype=np.double)
    for i in prange(nslopes):
        xr, slope = leaf_ranges[i], leaf_slopes[i]
        # Compute slope all the way across uniq_x but then trim line so
        # slope is only valid in range xr; don't set slope on right edge
        slopes[:, i] = np.where( (uniq_x < xr[0]) | (uniq_x >= xr[1]), np.nan, slope)

    # Slope values could be genuinely zero so we use nan not 0 for out-of-range.

    # Now average horiz across the matrix, averaging within each range
    # It's possible that some some rows would be purely NaN, indicating there are no
    # slopes for that X[colname] value. This can happen when we ignore some leaves,
    # when they have a single unique X[colname] value.

    # Compute:
    #   avg_slope_at_x = np.mean(slopes[good], axis=1)  (numba doesn't allow axis arg)
    #   slope_counts_at_x = nslopes - np.isnan(slopes).sum(axis=1)
    avg_slope_at_x = np.zeros(shape=nx)
    slope_counts_at_x = np.zeros(shape=nx)
    for i in prange(nx):
        row = slopes[i, :]
        n_nan = np.sum(np.isnan(row))
        avg_slope_at_x[i] = np.nan if n_nan==nslopes else np.nanmean(row)
        slope_counts_at_x[i] = nslopes - n_nan

    # return average slope at each unique x value and how many slopes included in avg at each x
    return avg_slope_at_x, slope_counts_at_x

# Hideous copying of avg_values_at_x_jit() to get different kinds of jit'ing. This is slower by 20%
# than other version but can run in parallel with multiprocessing package.
@jit(nopython=True)
def avg_slopes_at_x_nonparallel_jit(uniq_x, leaf_ranges, leaf_slopes):
    """
    Compute the weighted average of leaf_slopes at each uniq_x.

    Value at max(x) is NaN since we have no data beyond that point and so there is
    no forward difference.
    """
    nx = len(uniq_x)
    nslopes = len(leaf_slopes)
    slopes = np.zeros(shape=(nx, nslopes))
    for i in range(nslopes):
        xr, slope = leaf_ranges[i], leaf_slopes[i]
        # Compute slope all the way across uniq_x but then trim line so
        # slope is only valid in range xr; don't set slope on right edge
        slopes[:, i] = np.where( (uniq_x < xr[0]) | (uniq_x >= xr[1]), np.nan, slope)

    # Slope values could be genuinely zero so we use nan not 0 for out-of-range.

    # Now average horiz across the matrix, averaging within each range
    # It's possible that some some rows would be purely NaN, indicating there are no
    # slopes for that X[colname] value. This can happen when we ignore some leaves,
    # when they have a single unique X[colname] value.

    # Compute:
    #   avg_value_at_x = np.mean(slopes[good], axis=1)  (numba doesn't allow axis arg)
    #   slope_counts_at_x = nslopes - np.isnan(slopes).sum(axis=1)
    avg_value_at_x = np.zeros(shape=nx)
    slope_counts_at_x = np.zeros(shape=nx)
    for i in range(nx):
        row = slopes[i, :]
        n_nan = np.sum(np.isnan(row))
        avg_value_at_x[i] = np.nan if n_nan==nslopes else np.nanmean(row)
        slope_counts_at_x[i] = nslopes - n_nan

    # return average slope at each unique x value and how many slopes included in avg at each x
    return avg_value_at_x, slope_counts_at_x

def catwise_leaves(rf, X_not_col, X_col, y, max_catcode):
    """
    Return a 2D array with the average y value for each category in each leaf.
    Choose the cat code of smallest avg y as the reference category. I used to think it
    was arbitrary but now I realize that we have to mimic what StratPD does for
    numerical variables. By integrating the y deltas, we get back exactly the y
    values minus the minimum y. E.g., for x=[1,2,3] and y=[7,10,15], StratPD gets
    dx = [1,1] and dy=[3,5] then cumsum of dy gets [3,8] and we stick 0 on front:
    pdpy = [0,3,8].  To mimic the same thing with CatStratPD, we simply subtract the
    minimum for all categories in this leaf. If cats=[1,2,3] and y=[7,10,15], then
    we subtract 7 from y = [0,3,8] and we arrive at the same answer. For situations
    where the numerical var PD never skips over a valid x (extrapolating slope), it
    should give the same answer as categorical var PD. Subtracting the min y also
    has the benefit that debugging is easier because, within a leaf, all deltas are
    positive. During merging, we can still get negative deltas as we lineup cats/deltas,
    but those are true values. We should not shift everything to be positive for
    the final merged pdpy.

    Normalize the y values into deltas by subtracting the avg y value for the
    reference category from the avg y for all categories. The reference category is
    the category of the minimum avg y.

    The columns are the y avg value changes per cat found in a single leaf as
    they differ from the reference cat y average. Each row represents a category level. E.g.,

    row
    cat           leaf0       leaf1
     0       166.430176  186.796956
     1       219.590349  176.448626

    Cats are possibly noncontiguous with nan rows for cat codes not present. Not all
    values in a leaf column will be non-nan.  Only those categories mentioned in
    a leaf have values.  Shape is (max cat + 1, num leaves).

    Within a single leaf, there will typically only be a few categories represented.
    """
    leaves = leaf_samples(rf, X_not_col)

    leaf_deltas = np.full(shape=(max_catcode+1, len(leaves)), fill_value=np.nan)
    leaf_counts = np.zeros(shape=(max_catcode+1, len(leaves)), dtype=int)
    keep_leaf_idxs = np.full(shape=(len(leaves),), fill_value=True, dtype=bool)

    ignored = 0
    for leaf_i in range(len(leaves)):
        sample = leaves[leaf_i]
        leaf_cats = X_col[sample]
        leaf_y = y[sample]
        # perform a groupby(catname).mean()
        uniq_leaf_cats, count_leaf_cats = np.unique(leaf_cats, return_counts=True) # comes back sorted
        avg_y_per_cat = np.array([leaf_y[leaf_cats==cat].mean() for cat in uniq_leaf_cats])
        # print("uniq_leaf_cats",uniq_leaf_cats,"count_y_per_cat",count_leaf_cats)

        # if len(uniq_leaf_cats)==1 then we have single cat avg y and its delta is 0
        # but keep it as we'll treat as isolated group later and add marginal y for
        # this cat to get partial dependence value.

        # Can use any cat code as refcat; same "shape" of delta vec regardless of which we
        # pick. The vector is shifted/up or down but cat y's all still have the same relative
        # delta y. Might as well just pick the cat of smallest avg y.
        # Previously, I picked a random# reference category but that is unnecessary.
        # We will shift this vector during the merge operation so which we pick
        # here doesn't matter.
        idx_of_ref_cat_in_leaf = np.nanargmin(avg_y_per_cat)
        delta_y_per_cat = avg_y_per_cat - avg_y_per_cat[idx_of_ref_cat_in_leaf]
        # print("delta_y_per_cat",delta_y_per_cat)

        # Store into leaf i vector just those deltas we have data for
        # leave cats w/o representation as nan (uses index to figure out which rows to alter)
        leaf_deltas[uniq_leaf_cats, leaf_i] = delta_y_per_cat
        leaf_counts[uniq_leaf_cats, leaf_i] = count_leaf_cats

    # See unit test test_catwise_leaves:test_two_leaves_with_2nd_ignored()
    leaf_deltas = leaf_deltas[:,keep_leaf_idxs]
    leaf_counts = leaf_counts[:,keep_leaf_idxs]
    return leaf_deltas, leaf_counts, ignored

def cat_partial_dependence(X, y,
                           colname,  # X[colname] expected to be numeric codes
                           max_catcode=None,  # if we're bootstrapping, might see diff max's so normalize to one max
                           n_trees=1,
                           min_samples_leaf=5,
                           max_features=1.0,
                           rf_bootstrap=False,
                           supervised=True,
                           verbose=False):
    X_not_col = X.drop(colname, axis=1).values
    X_col = X[colname].values
    if (X_col<0).any():
        raise ValueError(f"Category codes must be > 0 in column {colname}")
    if not np.issubdtype(X_col.dtype, np.integer):
        raise ValueError(f"Category codes must be integers in column {colname} but is {X_col.dtype}")
    if max_catcode is None:
        max_catcode = np.max(X_col)
    if supervised:
        rf = RandomForestRegressor(n_estimators=n_trees,
                                   min_samples_leaf=min_samples_leaf,
                                   bootstrap = rf_bootstrap,
                                   max_features = max_features,
                                   oob_score=False)
    else:
        print("USING UNSUPERVISED MODE")
        X_synth, y_synth = conjure_twoclass(X)
        rf = RandomForestClassifier(n_estimators=n_trees,
                                    min_samples_leaf=min_samples_leaf,  # * 2, # there are 2x as many samples (X,X') so must double leaf size
                                    bootstrap=rf_bootstrap,
                                    max_features=max_features,
                                    oob_score=False)
        rf.fit(X_synth.drop(colname,axis=1), y_synth)

    rf.fit(X_not_col, y)
    if verbose and supervised:
        print(f"CatStrat Partition RF: dropping {colname} training R^2 {rf.score(X_not_col, y):.2f}")

    leaf_deltas, leaf_counts, ignored = \
        catwise_leaves(rf, X_not_col, X_col, y.values, max_catcode)

    uniq_x = np.unique(X_col)
    # Ignoring other vars, what is average y for all records with same catcode?
    # We need if avg_values_at_cat finds disjoint sets
    marginal_avg_y_per_cat = np.full(shape=(max_catcode+1,), fill_value=np.nan)
    for cat in uniq_x:
        marginal_avg_y_per_cat[cat] = y[X_col == cat].mean()

    avg_per_cat, count_per_cat = \
        avg_values_at_cat(leaf_deltas, leaf_counts, marginal_avg_y_per_cat, verbose=verbose)

    if verbose:
        print(f"CatStratPD Num samples ignored {ignored} for {colname}")

    return leaf_deltas, leaf_counts, avg_per_cat, count_per_cat, ignored

def avg_values_at_cat(leaf_deltas, leaf_counts,
                      marginal_avg_y_per_cat=None, # only used if disjoint cat sets exist
                      max_iter=3, verbose=False):
    """
    Merge the deltas for categories from all of the leaves. If there
    is at least one category in common, then to groups of deltas can
    be merged. It's possible, however, that there are multiple
    disjoint groups of categories with no cats in common. Previously I
    would just ignore those, but now I realize there could be
    important groups of categories that don't get included. For
    example with the bulldozer data set it was common to see the vast
    majority of categories being ignored. Obviously this dramatically
    underestimates the impact of a categorical variable.

    We don't have categories in common between disjoint groups but we
    can still estimate their relative position, otherwise we'd be
    forced to start all disjoint groups at zero. If we treat each
    group as a meta-category, then we can compare the average y for
    categories in a group. For meta-cats A, B, C with avg y =
    [100,130,40] then the meta-deltas are [0,30,-60].  The base for
    the zeros in A, B, C will be 0, 30, -60 so we just add that to all
    cat deltas in A, B, C.  Note: the y used for mean is not the leaf
    deltas but y from training set.

    This is effectively the marginal delta between those meta-groups,
    but the fact that they are disjoint means that there were no
    situations where, all else being equal, the associated leaves had
    categories in common. While the overall variable might be highly
    codependent, such as ModelID (with YearMade, MachineHours,...) on
    bulldozer price, the individual groups WITHIN a variable could be
    totally independent. Imagine typical construction bulldozers
    versus the giant earthmovers used in mining.  It's extremely
    unlikely that stratification would yield groups containing both
    construction bulldozers and mining bulldozers. That hints that
    the average difference in price between the two groups would not be
    biased (too much). Besides, it's the best we can do given we
    literally have no information about the relative cost of a mining
    bulldozer and a construction bulldozer. All we can do is compare
    the average price of a bulldozer in both categories, and use that
    to shift the two meta-cats relative to each other.

    Rather than a complicated data structure and algorithm to find all
    disjoint subsets of intersecting categories, it's easier to use the
    existing algorithm I had already that merged leaves with categories in
    common.  After merging all possible, there remains a set of leaf
    indexes indicating what leaves remain to process. So, just call function
    avg_values_at_cat_one_disjoint_region() repeatedly until the work
    set comes back as empty. That indicates all leaves have been
    integrated. The count per category vectors can simply be added
    together to get the final count per category. The catavg is the
    running vector average computed in each pass over the leaves. For
    that, we can add them vectors together to get a single vector
    (like a set union), but we still need to keep track of the different
    disjoint category groups. We can do this by collecting a list
    of cat sets.  Then we compute the average y for those categories
    for each group. Then update the catavg vector values for each
    group.

    See comment for avg_values_at_cat_one_disjoint_region()
    """
    # catavg is the running average vector and starts out as the first column (1st leaf's deltas)
    catavg = np.full(shape=(leaf_deltas.shape[0],), fill_value=np.nan, dtype=np.float)
    catgroups = []
    count_per_cat = np.zeros(shape=(leaf_deltas.shape[0],), dtype=np.int)
    work = range(0, leaf_deltas.shape[1])  # all leaf indexes added to work list
    # print("START work", work)
    while len(work)>0:
        catavg_, count_per_cat_, work = \
            avg_values_at_cat_one_disjoint_region(work, leaf_counts, leaf_deltas, max_iter, verbose)
        # print("catavg", catavg_)
        # print("count_per_cat", count_per_cat_)
        # print("remaining work",work)
        catgroup = np.where(count_per_cat_>0)[0]
        # print("catgroup", catgroup)
        catgroups.append(catgroup)
        catavg[catgroup] = catavg_[catgroup] # merge disjoint avgs for catgroup into running vector
        count_per_cat += count_per_cat_

    if len(catgroups)>1: # more than one disjoint set of categories?
        avgs = []
        for cats in catgroups:
            m = np.mean(marginal_avg_y_per_cat[cats])
            # print("cats",cats,"mean",m)
            avgs.append(m)
        # Do what StratPD does: all values are relative to left edge
        relative_group_y = np.array(avgs) - avgs[0]
        # print("group avgs", avgs)
        # print("relative_group_y", relative_group_y)
        for cats,group_deltay in zip(catgroups, relative_group_y):
            catavg[cats] += group_deltay
        # print("adjusted catavg", catavg)

    # if verbose: print("final cat avgs", parray3(catavg))
    return catavg, count_per_cat

def avg_values_at_cat_one_disjoint_region(work, leaf_counts, leaf_deltas, max_iter, verbose):
    """
    In leaf_deltas, we have information from the leaves indicating how
    much above or below each category was from the reference category
    of that leaf.  The reference category was arbitrarily selected but
    now I realize it must be the cat code associated with the minimum
    avg cat y.  All deltas within a leaf are therefore positive values
    and the relative y for the reference category is 0.  Categories
    not mentioned in the leaf, will have NAN values.

    The goal is to merge all of the columns in leaf_deltas, despite
    the fact that they do not have the same reference category. We
    init a running average vector to be the first column of category
    deltas. Then we attempt to merge each of the other columns into
    the running average. We make multiple passes over the columns of
    leaf_deltas until nothing changes, we hit the maximum number of
    iterations, or everything has merged.

    To merge vector v (column j of leaf_deltas) into catavg, select a
    category, index ix, in common at random.  Subtract v[ix] from v so
    that ix is v's new reference and v[ix]=0. Add catavg[ix] to the
    adjusted v so that v is now comparable to catavg. We can now do a
    weighted average of catavg and v, paying careful attention of NaN.

    Previously, I was picking a random category for merging in an
    effort to reduce the effect of outliers, with the assumption that
    outliers were rare. Given 5 categories in common between the
    running average vector and a new vector, randomly picking one
    means a 1/5 chance of picking the outlier.  Outliers as reference
    categories shift the outlierness to all other categories. Boooo

    Now, I select the category in common that has the most evidence:
    the category associated with the most number of observations.

    It's possible that more than a single value within a leaf_deltas
    vector is 0.  I.e., the reference category value is always 0 in
    the vector, but there might be another category whose value was
    the same y, giving a 0 relative value.

    Example:

    leaf_deltas

    [[ 0. nan nan]
     [ 1. nan nan]
     [nan  0. nan]
     [nan  3.  0.]
     [nan  3.  2.]
     [ 4. nan nan]
     [ 5. nan nan]]

    leaf_counts

    [[1 0 0]
     [1 0 0]
     [0 1 0]
     [0 1 1]
     [0 1 1]
     [1 0 0]
     [1 0 0]]
    """
    initial_leaf_idx = work[0]
    work = set(work[1:])
    catavg = leaf_deltas[:, initial_leaf_idx]  # init with first ref category (column)
    catavg_weight = leaf_counts[:, initial_leaf_idx]
    completed = {-1}  # init to any nonempty set to enter loop
    iteration = 1
    # Three passes should be sufficient to merge all possible vectors, but
    # I'm being paranoid here and allowing it to run until completion or some maximum iterations
    while len(work) > 0 and len(completed) > 0 and iteration <= max_iter:
        # print(f"PASS {iteration} len(work)", len(work))
        completed = set()
        for j in work:  # for remaining leaf index in work list, avg in the vectors
            v = leaf_deltas[:, j]
            are_intersecting = ~np.isnan(catavg) & ~np.isnan(v)
            intersection_idx = np.where(are_intersecting)[0]

            # print(intersection_idx)
            if len(intersection_idx) == 0:  # found something to merge into catavg?
                continue

            # Merge column j into catavg vector
            # cat for merging is the one with most supporting evidence
            cur_weight = leaf_counts[:, j]

            PICK_MOST_CONF = False
            if PICK_MOST_CONF:
                ix = np.argmax(np.where((cur_weight > 0) & are_intersecting, cur_weight, 0))
                shifted_v = v - v[ix]  # make ix the reference cat in common
                relative_to_value = catavg[ix]  # corresponding value in catavg
                adjusted_v = shifted_v + relative_to_value  # adjust so v is mergeable with catavg
            else:
                adjusted_v = compute_avg_merge_candidate(catavg, v, intersection_idx)

            prev_catavg = catavg  # track only for verbose/debugging purposes
            catavg = nanavg_vectors(catavg, adjusted_v, catavg_weight, cur_weight)
            # Update weight of running avg to incorporate "mass" from v
            catavg_weight += cur_weight
            if verbose:
                # TODO: broken
                print(f"{ix:-2d} : vec to add =", parray(v), f"- {v[ix]:.2f}")
                print("     shifted    =", parray(shifted_v),
                      f"+ {relative_to_value:.2f}")
                print("     adjusted   =", parray(adjusted_v), "*", cur_weight)
                print("     prev avg   =", parray(prev_catavg), "*",
                      catavg_weight - cur_weight)
                print("     new avg    =", parray(catavg))
                print()
            completed.add(j)
        iteration += 1
        work = work - completed
    return catavg, catavg_weight, list(work)

def compute_avg_merge_candidate(catavg, v, intersection_idx):
    """
    Given intersecting deltas in catavg and v, compute average delta
    one could merge into running average.  If one cat is an outlier,
    picking that really distorts the vector we merge into running
    average vector. So, effectively merge using all as the ref
    cat in common by merging in average of all possible refcats.

    When there is no noise in y, the average merge candidate is
    the same as any single candidate. So, with no noise, we get
    exact answer; averaging here doesn't cost us anything. It
    only helps to spread noise across categories.
    """
    merge_candidates = []
    for i in intersection_idx:
        merge_candidates.append(v - v[i] + catavg[i])
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)
        # We get "Mean of empty slice" when all entries are Nan but we want that.
        v = np.nanmean(merge_candidates, axis=0)
    return v

# -------------- S U P P O R T ---------------

def scramble(X : np.ndarray) -> np.ndarray:
    """
    From Breiman: https://www.stat.berkeley.edu/~breiman/RandomForests/cc_home.htm
    "...the first coordinate is sampled from the N values {x(1,n)}. The second
    coordinate is sampled independently from the N values {x(2,n)}, and so forth."
    """
    X_rand = X.copy()
    ncols = X.shape[1]
    for col in range(ncols):
        X_rand[:,col] = np.random.choice(X[:,col], len(X), replace=True)
    return X_rand

def df_scramble(X : pd.DataFrame) -> pd.DataFrame:
    """
    From Breiman: https://www.stat.berkeley.edu/~breiman/RandomForests/cc_home.htm
    "...the first coordinate is sampled from the N values {x(1,n)}. The second
    coordinate is sampled independently from the N values {x(2,n)}, and so forth."
    """
    X_rand = X.copy()
    for colname in X:
        # X_rand[colname] = np.random.choice(X[colname], len(X), replace=True)
        X_rand[colname] = X_rand[colname].sample(frac=1.0)
    return X_rand

def conjure_twoclass(X):
    """
    Make new data set 2x as big with X and scrambled version of it that
    destroys structure between features. Old is class 0, scrambled is class 1.
    """
    if isinstance(X, pd.DataFrame):
        X_rand = df_scramble(X)
        X_synth = pd.concat([X, X_rand], axis=0)
    else:
        X_rand = scramble(X)
        X_synth = np.concatenate([X, X_rand], axis=0)
    y_synth = np.concatenate([np.zeros(len(X)),
                              np.ones(len(X_rand))], axis=0)
    return X_synth, pd.Series(y_synth)

def nanavg_vectors(a, b, wa=1.0, wb=1.0):
    """
    Add two vectors a+b but support nan+x==x and nan+nan=nan
    np.nanmean works to get nan+nan=nan, but for weighted avg
    we need to divide by wa+wb after using nansum. nansum gives
    0 not nan it seems when adding nan+nan. Do it the hard way.
    """
    a_nan = np.isnan(a)
    b_nan = np.isnan(b)
    c = a*wa + b*wb               # weighted average where both are non-nan
    c /= zero_as_one(wa+wb)       # weighted avg
    c[a_nan] = b[a_nan]           # copy any stuff where b has only value (unweighted into result)
    in_a_not_b = (~a_nan) & b_nan
    c[in_a_not_b] = a[in_a_not_b] # copy stuff where a has only value
    return c

def nanmerge_matrix_cols(A):
    """
    Add all vertical vectors in A but support nan+x==x and nan+nan=nan.
    """
    s = np.nansum(A, axis=1)
    all_nan_entries = np.isnan(A)
    # if all entries for a cat are nan, make sure sum s is nan for that cat
    s[all_nan_entries.all(axis=1)] = np.nan
    return s

def zero_as_one(a):
    return np.where(a == 0, 1, a)

def parray(a):
    if type(a[0])==np.int64:
        return '[ ' + (' '.join([f"{x:6d}" for x in a])).strip() + ' ]'
    else:
        return '[ ' + (' '.join([f"{x:6.2f}" for x in a])).strip() + ' ]'

def parray3(a):
    return '[ ' + (' '.join([f"{x:6.3f}" for x in a])).strip() + ' ]'


from enum import Enum

from pyspark import SparkContext

from geospark.core.SpatialRDD import SpatialRDD, PolygonRDD, LineStringRDD, PointRDD
from geospark.utils.decorators import require

class DiscLoader:

    @classmethod
    def load(cls, sc: SparkContext, path: str) -> SpatialRDD:
        raise NotImplementedError()

class IndexDiscLoader:

    @classmethod
    def load(cls, sc: SparkContext, path: str):
        jvm = sc._jvm
        index_rdd = jvm.ObjectSpatialRDDLoader.loadIndexRDD(sc._jsc, path)
        return index_rdd

class PolygonRDDDiscLoader(DiscLoader):

    @classmethod
    @require(["ObjectSpatialRDDLoader"])
    def load(cls, sc: SparkContext, path: str) -> SpatialRDD:
        jvm = sc._jvm
        polygon_rdd = PolygonRDD()
        srdd = jvm.ObjectSpatialRDDLoader.loadPointSpatialRDD(sc._jsc, path)
        polygon_rdd.set_srdd(srdd)
        return polygon_rdd

class PointRDDDiscLoader(DiscLoader):

    @classmethod
    @require(["ObjectSpatialRDDLoader"])
    def load(cls, sc: SparkContext, path: str) -> SpatialRDD:
        jvm = sc._jvm
        polygon_rdd = PointRDD()
        srdd = jvm.ObjectSpatialRDDLoader.loadPointSpatialRDD(sc._jsc, path)
        polygon_rdd.set_srdd(srdd)
        return polygon_rdd

class LineStringRDDDiscLoader(DiscLoader):

    @classmethod
    @require(["ObjectSpatialRDDLoader"])
    def load(cls, sc: SparkContext, path: str) -> SpatialRDD:
        jvm = sc._jvm
        line_string_rdd = LineStringRDD()
        srdd = jvm.ObjectSpatialRDDLoader.loadLineStringSpatialRDD(sc._jsc, path)
        line_string_rdd.set_srdd(srdd)
        return line_string_rdd

class SpatialRDDDiscLoader(DiscLoader):

    @classmethod
    @require(["ObjectSpatialRDDLoader"])
    def load(cls, sc: SparkContext, path: str) -> SpatialRDD:
        jvm = sc._jvm
        spatial_rdd = SpatialRDD()
        srdd = jvm.ObjectSpatialRDDLoader.loadSpatialRDD(sc._jsc, path)
        spatial_rdd.set_srdd(srdd)
        return spatial_rdd

class GeoType(Enum):
    POINT = "POINT"
    POLYGON = "POLYGON"
    LINESTRING = "LINESTRING"
    GEOMETRY = "GEOMETRY"

loaders = {
    GeoType.POINT: PointRDDDiscLoader,
    GeoType.POLYGON: PolygonRDDDiscLoader,
    GeoType.LINESTRING: LineStringRDDDiscLoader,
    GeoType.GEOMETRY: SpatialRDDDiscLoader
}

def load_spatial_rdd_from_disc(sc: SparkContext, path: str, geometry_type: GeoType):
    """

    :param sc:
    :param path:
    :param geometry_type:
    :return:
    """
    return loaders[geometry_type].load(sc, path)

def load_spatial_index_rdd_from_disc(sc: SparkContext, path: str):

    return IndexDiscLoader.load(sc, path)


# -*- coding: utf-8 -*-
"""
Machine Learning based Analytical Continuation

@author: Junyin Zhang
"""
import scipy.io as sio
import numpy as np
import tensorflow as tf
matfn = u'C:/Users/ienor/Desktop/Hubbard_Advanced/MLAC/' + 'U0.mat'
batch = sio.loadmat(matfn)['mea_ML_data_p']
#mea_ML_data = NumOfVertexs*NumOfVertexs*TempSlice
#px = (px_index-NumInEdge) * delta_p - pi
NumInEdge = np.shape(batch)[1]
TempSlice = np.shape(batch)[2]

kx_index = 1
ky_index = 1
Beta = 2.0
G_k = np.zeros([1,TempSlice])
G_k[0] = batch[kx_index,ky_index,0]
for time_index in range(TempSlice-1):
    G_k[0,time_index+1] = batch[kx_index,ky_index,time_index]

Gauge = np.zeros([TempSlice,64])
for t_index in range(TempSlice):
    for e_index in range(64):
        tau = t_index*Beta/TempSlice
        omega = -3.2 + 0.1*e_index
        val = np.exp(-omega*(tau-Beta/2))/(2*np.cosh(Beta*omega/2))
        Gauge[t_index,e_index] = val

Gauge_tf = tf.convert_to_tensor(Gauge, dtype=tf.float32);
#################################################################
import keras
from keras import models
from keras import layers
from keras.layers import Input,Dense
from keras import backend as K

num_epochs = 30;
num_train = 1024;
num_test = 1e1;
train_data = np.zeros([round(num_train),TempSlice])
test_data = np.zeros([round(num_test),TempSlice])
for index in range(round(num_train)):
    train_data[index,:] = G_k;
for index in range(round(num_test)):
    test_data[index,:] = G_k

def custom_loss_wrapper(input_tensor):
    def custom_loss(y_true,y_pred):
        #G_pred = K.dot(Gauge_tf,y_pred)
        loss_value = K.sum(K.square(K.dot(Gauge_tf,K.transpose(y_pred))-K.transpose(input_tensor)))
        return loss_value
    return custom_loss

def baseline_model():
    input_tensor = Input(shape=(16,))
    hidden_1 = Dense(64, activation='sigmoid')(input_tensor)
    hidden_2 = Dense(32, activation='relu')(hidden_1)
    out = Dense(64, activation='sigmoid')(hidden_2)
    model = keras.Model(input_tensor, out)
    model.compile(optimizer='rmsprop',
                loss = custom_loss_wrapper(input_tensor))
    return model

model = baseline_model()

history = model.fit(train_data,train_data,
                    epochs = num_epochs,batch_size = 1,verbose=1)
#history = history.history['val_mean_absolute_error']

A_pred = np.transpose(model.predict(G_k))
G_pred = np.transpose(np.dot(Gauge,A_pred))
G_diff_r = (G_pred - G_k)/G_k
print(np.linalg.norm(G_diff_r))

#sample_tf = tf.convert_to_tensor(sample, dtype=tf.float32);