vietnam = ee.FeatureCollection('users/servirmekong/countries/VNM_adm1')
compositeSentinel2 = ee.Image('users/phamngochairs/NB_S2A_Spectral_2018_2018')

# export training data
aquaculture =  ee.FeatureCollection("users/servirmekong/referenceData/aquaculture");
barren = ee.FeatureCollection("users/servirmekong/referenceData/barren");
cropland = ee.FeatureCollection("users/servirmekong/referenceData/cropland");
deciduousForest = ee.FeatureCollection("users/servirmekong/referenceData/deciduousForest");
evergreen = ee.FeatureCollection("users/servirmekong/referenceData/evergreen_forest");
floodedForest = ee.FeatureCollection("users/servirmekong/referenceData/floodedForest");
plantations = ee.FeatureCollection("users/servirmekong/referenceData/plantations");
grass = ee.FeatureCollection("users/servirmekong/referenceData/grass");
mangroves = ee.FeatureCollection("users/servirmekong/referenceData/mangrove");
mixedForest = ee.FeatureCollection("users/servirmekong/referenceData/mixedForest");
plantations = ee.FeatureCollection("users/servirmekong/referenceData/plantations");
rice = ee.FeatureCollection("users/servirmekong/referenceData/rice");
shrub = ee.FeatureCollection("users/servirmekong/referenceData/shrub");
snow = ee.FeatureCollection("users/servirmekong/referenceData/snow");
water = ee.FeatureCollection("users/servirmekong/referenceData/water");
imperv = ee.FeatureCollection("users/servirmekong/referenceData/imperv");
wetlands = ee.FeatureCollection("users/servirmekong/referenceData/wetlands");
noRice = aquaculture.merge(barren).merge(cropland).merge(deciduousForest).merge(evergreen).merge(floodedForest).merge(plantations).merge(grass).merge(mangroves).merge(mixedForest).merge(shrub).merge(snow).merge(water).merge(imperv).merge(wetlands);

deltaProvience = ['Dong Thap', 'Tien Giang', 'Vinh Long', 'Tra Vinh', 'An Giang', 'Can Tho', 'Hau Giang', 'Soc Trang', 'Bac Lieu', 'Kien Giang', 'Ben Tre', 'Long An'];
geometry = vietnam.filter(ee.Filter.inList('VARNAME_1', deltaProvience));

def _set_true_feature(feature):
    return feature.set(RESPONSE, 1)

def _set_false_feature(feature):
    return feature.set(RESPONSE, 0)

rice = rice.filterBounds(geometry).map(_set_true_feature).randomColumn();
trainingRice = rice.filterMetadata('random', 'greater_than', 0.3);
testingRice = rice.filterMetadata('random', 'less_than', 0.3);

noRice = noRice.filterBounds(geometry).map(_set_false_feature).randomColumn();
trainingNoRice = noRice.filterMetadata('random', 'greater_than', 0.3);
testingNoRice = noRice.filterMetadata('random', 'less_than', 0.3);


compositeSentinel2 = compositeSentinel2.clip(geometry).neighborhoodToArray(kernel);

trainingRiceSampled = compositeSentinel2.sampleRegions(
  collection = trainingRice,
  properties = [RESPONSE],
  scale = 10
)
trainingRiceList = trainingRice.toList(trainingRice.size());


testingRiceSampled = compositeSentinel2.sampleRegions(
  collection = testingRice,
  properties = [RESPONSE],
  scale = 10
)
testingRiceList = testingRice.toList(testingRice.size());

trainingNoRiceSampled = compositeSentinel2.sampleRegions(
  collection = trainingNoRice,
  properties = [RESPONSE],
  scale = 10
)
trainingNoRiceList = trainingNoRice.toList(trainingNoRice.size());

testingNoRiceSampled = compositeSentinel2.sampleRegions(
  collection = testingNoRice,
  properties = [RESPONSE],
  scale = 10
)
testingNoRiceList = testingNoRice.toList(testingNoRice.size());


# //////////////////////////////
# //Export all the training data (in many pieces), with one task per geometry.

# // These numbers determined experimentally.
n = 100; # Number of shards in each polygon.
N = 2000; # Total sample size in each polygon.


#var selectors = Array('blue','green','red','nir','swir1','swir2','rice');
selectors = compositeSentinel2.bandNames().getInfo();
selectors.append(RESPONSE)

# trainingRice testingRice trainingNoRice testingNoRice
# 11            6            11             6

#for g in range(trainingRice.size().getInfo()):
states = {
    'trainingRice': {
        'm': 11,
        'list': trainingRiceList
    },
    'trainingNoRice': {
        'm': 11,
        'list': trainingNoRiceList
    },
    'testingRice': {
        'm': 6,
        'list': testingRiceList
    },
    'testingNoRice': {
        'm': 6,
        'list': testingNoRiceList
    }
}

for state in states:
  TRAINING_BASE = state
  m = states[state]['m']
  _list = states[state]['list']
  for g in range(m):
    geomSample = ee.FeatureCollection([])
    for i in range(n):
      feature = ee.Feature(_list.get(g))
      sample = compositeSentinel2.sample(
        region = feature.geometry(),
        scale = 10,
        numPixels = N / n, # Size of the shard.
        seed = i,
        tileScale = 8
      ).map(_set_false_feature)
      geomSample = geomSample.merge(sample)

    desc = TRAINING_BASE + '_g' + str(g)
    task = ee.batch.Export.table.toCloudStorage(
      collection = geomSample,
      description = desc,
      bucket = BUCKET,
      fileNamePrefix = desc,
      fileFormat = 'TFRecord',
      selectors = selectors
    )
    task.start()